code: def sample_categorical(
      rng: jnp.ndarray,
      logits: jnp.ndarray,
      num_samples: int,
      *,
      replacement: bool = True,
  ):
      """Sample catogorical with or without replacement for the top-k selector.
  
      Args:
        rng: JAX PRNG key.
        logits: Categorical distribution logits of shape [batch_dims, num_classes].
        num_samples: Number of samples to produce.
        replacement: If True, sampling is done with replacement.
  
      Returns:
        Categorial samples of shape [batch_dims, num_samples].
      """
      rng = jax.random.split(rng, num_samples)
      if replacement:
          samples = jax.vmap(jax.random.categorical, in_axes=(0, None))(
              rng, logits
          )
      else:
          num_categories = logits.shape[-1]
          if num_categories < num_samples:
              raise ValueError(
                  f"Number of samples ({num_samples}) must be <= number of"
                  f" categories ({num_categories}) when sampling without"
                  f" replacement."
              )
  
          def sample_one(logits, scan_rng):
              samples = jax.random.categorical(
                  scan_rng, logits, axis=-1
              )
              mask = jax.nn.one_hot(
                  samples, num_categories, dtype=jnp.bool_
              )
              logits = jnp.where(mask, -1e10, logits)
              return logits, samples
  
          _, samples = jax.lax.scan(sample_one, logits, rng)
  
      # Restore original shape.
      ndim = samples.ndim
      if ndim > 1:
          samples = jnp.transpose(
              samples, axes=tuple(range(1, ndim)) + (0,)
          )
      return samples
  
imports: import abc
  import enum
  import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Iterable
  from absl import logging
  from flax import linen
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import attention_layers
label_desc: Use image node_id 5 for calling a global function with example usage: sample_categorical(rng, logits, num_samples) and returns: samples, logits,  samples