def categorical(logits, num_samples, dtype="int32", seed=None):
    """Draws samples from a categorical distribution.

    This function takes as input `logits`, a 2-D input tensor with shape
    (batch_size, num_classes). Each row of the input represents a categorical
    distribution, with each column index containing the log-probability for a
    given class.

    The function will output a 2-D tensor with shape (batch_size, num_samples),
    where each row contains samples from the corresponding row in `logits`.
    Each column index contains an independent samples drawn from the input
    distribution.

    Args:
        logits: 2-D Tensor with shape (batch_size, num_classes). Each row
            should define a categorical distribution with the unnormalized
            log-probabilities for all classes.
        num_samples: Int, the number of independent samples to draw for each
            row of the input. This will be the second dimension of the output
            tensor's shape.
        dtype: Optional dtype of the output tensor.
        seed: Optional Python integer or instance of
           `keras.random.SeedGenerator`.
            By default, the `seed` argument is `None`, and an internal global
            `keras.random.SeedGenerator` is used. The `seed` argument can be
            used to ensure deterministic (repeatable) random number generation.
            Note that passing an integer as the `seed` value will produce the
            same random values for each call. To generate different random
            values for repeated calls, an instance of
            `keras.random.SeedGenerator` must be provided as the `seed` value.
            Remark concerning the JAX backend: When tracing functions with the
            JAX backend the global `keras.random.SeedGenerator` is not
            supported. Therefore, during tracing the default value seed=None
            will produce an error, and a `seed` argument must be provided.

    Returns:
        A 2-D tensor with (batch_size, num_samples).
    """
    logits_shape = list(backend.convert_to_tensor(logits).shape)
    if len(logits_shape) != 2:
        raise ValueError(
            "`logits` should be a 2-D tensor with shape "
            f"[batch_size, num_classes]. Received: logits={logits}"
        )
    return backend.random.categorical(
        logits, num_samples, dtype=dtype, seed=seed
    )
Draws samples from a categorical distribution.

    This function takes as input `logits`, a 2-D input tensor with shape
    (batch_size, num_classes). Each row of the input represents a categorical
    distribution, with each column index containing the log-probability for a
    given class.

    The function will output a 2-D tensor with shape (batch_size, num_samples),
    where each row contains samples from the corresponding row in `logits`.
    Each column index contains an independent samples drawn from the input
    distribution.

    Args:
        logits: 2-D Tensor with shape (batch_size, num_classes). Each row
            should define a categorical distribution with the unnormalized
            log-probabilities for all classes.
        num_samples: Int, the number of independent samples to draw for each
            row of the input. This will be the second dimension of the output
            tensor's shape.
        dtype: Optional dtype of the output tensor.
        seed: Optional Python integer or instance of
           `keras.random.SeedGenerator`.
            By default, the `seed` argument is `None`, and an internal global
            `keras.random.SeedGenerator` is used. The `seed` argument can be
            used to ensure deterministic (repeatable) random number generation.
            Note that passing an integer as the `seed` value will produce the
            same random values for each call. To generate different random
            values for repeated calls, an instance of
            `keras.random.SeedGenerator` must be provided as the `seed` value.
            Remark concerning the JAX backend: When tracing functions with the
            JAX backend the global `keras.random.SeedGenerator` is not
            supported. Therefore, during tracing the default value seed=None
            will produce an error, and a `seed` argument must be provided.

    Returns:
        A 2-D tensor with (batch_size, num_samples).
    
keras/src/random/random.py