code: def random_categorical(logits, num_samples, seed):
      """Returns a sample from a categorical distribution. `logits` must be 2D."""
      probs = jax.nn.softmax(logits)
      cum_sum = jnp.cumsum(probs, axis=-1)
  
      eta = random.uniform(
          make_tensor_seed(seed), (num_samples,) + cum_sum.shape[:-1]
      )
      cum_sum = jnp.broadcast_to(
          cum_sum, (num_samples,) + cum_sum.shape
      )
  
      flat_cum_sum = cum_sum.reshape([-1, cum_sum.shape[-1]])
      flat_eta = eta.reshape([-1])
      return (
          jax.vmap(_searchsorted)(flat_cum_sum, flat_eta)
          .reshape(eta.shape)
          .T
      )
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  from jax import tree_util
  import jax.numpy
label_desc: Use image node_id 14 for calling a global function with example usage: random_categorical(logits, num_samples, seed) and returns: jax