jax.random.categorical#
- jax.random.categorical(key, logits, axis=-1, shape=None)[source]#
Sample random values from categorical distributions.
- Parameters:
key (KeyArrayLike) – a PRNG key used as the random key.
logits (RealArray) – Unnormalized log probabilities of the categorical distribution(s) to sample from, so that softmax(logits, axis) gives the corresponding probabilities.
axis (int) – Axis along which logits belong to the same categorical distribution.
shape (Shape | None) – Optional, a tuple of nonnegative integers representing the result shape. Must be broadcast-compatible with
np.delete(logits.shape, axis)
. The default (None) produces a result shape equal tonp.delete(logits.shape, axis)
.
- Return type:
Array
- Returns:
A random array with int dtype and shape given by
shape
ifshape
is not None, or elsenp.delete(logits.shape, axis)
.