jax.random.choiceΒΆ

jax.random.choice(key, a, shape=(), replace=True, p=None)[source]ΒΆ

Generates a random sample from a given 1-D array.

Parameters
  • key (ndarray) – a PRNGKey used as the random key.

  • a (Any) – 1D array or int. If an ndarray, a random sample is generated from its elements. If an int, the random sample is generated as if a were arange(a).

  • shape (Sequence[int]) – tuple of ints, optional. Output shape. If the given shape is, e.g., (m, n), then m * n samples are drawn. Default is (), in which case a single value is returned.

  • replace (bool) – boolean. Whether the sample is with or without replacement. default is True.

  • p – 1-D array-like, The probabilities associated with each entry in a. If not given the sample assumes a uniform distribution over all entries in a.

Return type

ndarray

Returns

An array of shape shape containing samples from a.