Introduction

tf.random.categorical samples class indices from one or more categorical distributions. The function is small, but the parameters matter because the input must be logits, the output shape depends on num_samples, and each row of the input tensor is treated as a separate distribution.

The core parameters

The most important arguments are:

• 'logits: a 2D tensor shaped like batch_size x num_classes'
• 'num_samples: how many draws to take from each row'
• 'dtype: integer type of the returned class IDs'
• 'seed: optional seed for reproducibility'

The function does not expect probabilities directly. It expects logits, which are unnormalized scores.

A tiny working example

Here are two categorical distributions, each with three classes:

1import tensorflow as tf
2
3logits = tf.math.log([
4    [0.1, 0.7, 0.2],
5    [0.8, 0.1, 0.1],
6])
7
8samples = tf.random.categorical(
9    logits=logits,
10    num_samples=5,
11    dtype=tf.int32,
12    seed=7,
13)
14
15print(samples)

The output shape is 2 x 5 because there are two rows in logits and five samples requested per row.

How to read the rows

Each row of logits is one distribution:

• first row favors class 1
• second row favors class 0

That means most samples from the first row should be 1, and most samples from the second row should be 0. The function does not mix the rows together. Sampling happens row by row.

Why logits are used

TensorFlow uses logits because many models already produce raw scores before a softmax layer. Passing logits directly avoids unnecessary conversions and can be more numerically stable.

If you already have probabilities, convert them carefully:

1import tensorflow as tf
2
3probabilities = tf.constant([[0.2, 0.5, 0.3]], dtype=tf.float32)
4logits = tf.math.log(probabilities)
5samples = tf.random.categorical(logits, num_samples=3)
6print(samples)

Passing raw probabilities as if they were logits changes the distribution, which is a common mistake.

What num_samples really means

num_samples is not the number of classes and not the batch size. It is the number of draws from each categorical distribution.

For example:

1import tensorflow as tf
2
3logits = tf.math.log([[0.4, 0.6]])
4samples = tf.random.categorical(logits, num_samples=10)
5print(samples.shape)

The shape is 1 x 10 because there is one distribution and ten draws.

dtype and seed

The returned samples are class IDs, so the output dtype is an integer type such as tf.int64 or tf.int32.

The seed parameter helps when you want repeatable experiments or tests. It does not make the result "more correct." It just makes the randomness easier to reproduce.

Shape mistakes are very common

Even when you only want to sample from one distribution, logits still needs to be 2D. That means a single distribution should look like [[...]], not just [...].

For example:

single_distribution = tf.math.log([[0.2, 0.5, 0.3]])
samples = tf.random.categorical(single_distribution, num_samples=4)
print(samples.shape)

That shape rule explains a lot of first-time confusion with this API, especially in quick experiments and notebooks for beginners who expect simpler inputs. It is a small but important detail.

Common Pitfalls

• Passing probabilities directly instead of logits.
• Forgetting that logits must be 2D, even when you have only one distribution.
• Misreading num_samples as the number of classes.
• Expecting the output values to be probabilities instead of sampled class indices.

Summary

• 'tf.random.categorical samples class IDs from each row of a logits tensor.'
• 'logits must be shaped like batch_size x num_classes.'
• 'num_samples controls how many draws are taken per row.'
• Use seed for reproducibility and dtype for the integer output type you want.