Skip to content

← JAX Stochasticity step 4 of 25

Medium Primitives

Categorical Sampling

Jump to the task

Draw a class index from scores

A categorical distribution chooses one of K classes. Its logits are scores whose softmax gives the class probabilities. Larger scores mean more likely classes, but sampling need not choose the largest score every time. That would be argmax, not a random draw.

For example, logits [0, log(3)] describe probabilities [0.25, 0.75]. Adding the same constant to both scores leaves these probabilities unchanged. If you already have positive probabilities, their logarithms are suitable logits. Passing probabilities directly as logits changes the distribution: softmax([0.25, 0.75]) is about [0.378, 0.622].

Sample indices, not probabilities

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
logits = jnp.array([1., 1., 1.])
samples = jax.random.categorical(key, logits, shape=(5,))
# [1, 1, 2, 1, 2] with the exercise runtime's default settings
result = samples.astype(jnp.float32)
# [1., 1., 2., 1., 2.]

Equal logits make the three classes equally likely. They do not force a small sample to contain equal counts of each class. Sampling is with replacement here, so repeated indices are expected and n can exceed K.

The output contains class positions 0 through K - 1. It is neither a vector of probabilities nor a one-hot matrix. For this one-dimensional logits input, shape (n,) requests n draws from the same distribution. With batched logits, output shape must also accommodate the distribution batch dimensions; the simple sample-count rule is specific to this task.

Casting does not make a discrete draw differentiable

Categorical sampling returns integer indices. This exercise asks for float32 output, so cast the result. That is an output-format choice, not a way to obtain useful pathwise gradients through the class choice. The selection is discrete; softmax probabilities being differentiable does not make sampled indices differentiable in the same way.

You do not need to compute softmax before calling categorical. The probabilities are defined by softmax(logits); that does not mean the sampler literally computes a softmax array internally.

Problem

Implement categorical_sample(seed, logits, n) using one call to jax.random.categorical with shape (int(n),), then cast to float32.

  • Seed is a whole Python number in [0, 2**32), supplied as a float. Create the key with jax.random.PRNGKey(int(seed)).
  • logits is a nonempty one-dimensional float32 array of finite scores.
  • n is a nonnegative whole Python number supplied as a float.
  • Return shape (int(n),) in the original draw order. Zero draws returns an empty vector; one class always gives index zero.

Use the default runtime sampler and Threefry configuration, with the original seed-derived key. Splitting it first changes the expected samples. Repeating the same inputs and settings reproduces the output.

For JIT, mark seed and n static: jax.jit(categorical_sample, static_argnames=("seed", "n")). Logit values can remain dynamic. int() cannot make a traced argument static.