keras/src/backend/numpy/random.py/categorical
def categorical(logits, num_samples, dtype="int64", seed=None):
    seed = draw_seed(seed)
    rng = np.random.default_rng(seed)
    output = []
    for logits_instance in logits:
        probabilities = softmax(logits_instance)
        classes = np.arange(logits_instance.shape[-1])
        samples = rng.choice(classes, size=num_samples, p=probabilities)
        output.append(samples)
    return np.array(output).astype(dtype)