keras/src/backend/jax/random.py/categorical
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)