code: def categorical(logits, num_samples, dtype="int32", seed=None):
      seed = jax_draw_seed(seed)
      output_shape = list(logits.shape)
      output_shape[1] = num_samples
      output_shape = tuple(output_shape)
      output = jax.random.categorical(
          seed, logits[..., None], shape=output_shape, axis=1
      )
      return output.astype(dtype)
  
file_path: keras/keras/backend/jax/random.py
imports: import jax
  from keras.backend.config import floatx
  from keras.random.seed_generator import SeedGenerator
  from keras.random.seed_generator import draw_seed
  from keras.random.seed_generator import make_default_seed
label_desc: Use image node_id 4 for calling a global function with example usage: categorical(logits, num_samples, dtype, seed) and returns: output