jax.random.categorical#

jax.random.categorical(key, logits, axis=-1, shape=None)[source]#

Sample random values from categorical distributions.

Parameters

  • key (Union[Array, PRNGKeyArray]) – a PRNG key used as the random key.

  • logits (Union[Array, ndarray, bool_, number, bool, int, float, complex]) – 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 (Optional[Sequence[int]]) – 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 to np.delete(logits.shape, axis).

Return type

Array

Returns

A random array with int dtype and shape given by shape if shape is not None, or else np.delete(logits.shape, axis).