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 – a PRNGKey used as the random key.
a – 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 – 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 – 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.

Returns

An array of shape shape containing samples from a.