code: def _categorical_jax(
      logits, num_samples, dtype=None, seed=None, name=None
  ):  # pylint: disable=unused-argument
      """Jax implementation of `tf.random.stateless_categorical`."""
      dtype = utils.numpy_dtype(dtype or np.int64)
      if not hasattr(logits, "shape") or not hasattr(logits, "dtype"):
          logits = np.array(logits, np.float32)
      import jax.random as jaxrand  # pylint: disable=g-import-not-at-top
  
      if seed is None:
          raise ValueError("Must provide PRNGKey to sample in JAX.")
      z = jaxrand.gumbel(
          key=seed,
          shape=logits.shape + (num_samples,),
          dtype=logits.dtype,
      )
      return np.argmax(np.expand_dims(logits, -1) + z, axis=-2).astype(
          dtype
      )
  
imports: import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softmax
label_desc: Use image node_id 5 for calling a global function with example usage: _categorical_jax(logits, num_samples, dtype, seed, name) and returns: np