keras/src/backend/tensorflow/random.py/categorical
def categorical(logits, num_samples, dtype="int64", seed=None):
    seed = _cast_seed(draw_seed(seed))
    output = tf.random.stateless_categorical(logits, num_samples, seed=seed)
    return tf.cast(output, dtype)