Skip to content

← JAX Stochasticity step 11 of 25

Medium Primitives

Gumbel-Max Categorical Sampling

Jump to the task

Why this matters

To sample a class, add independent standard Gumbel noise to each logit and choose the largest resulting score. In exact arithmetic, this Gumbel-Max construction chooses class i with probability softmax(logits)[i]. It turns a weighted random choice into noise generation followed by a maximum search.

The output is a hard class index. Replacing that maximum selection with a temperature-scaled softmax gives the Gumbel-Softmax relaxation, but that is a different output and a different exercise.

Worked mini-example

import jax
import jax.numpy as jnp

logits = jnp.array([1.0, 2.0, 3.0])
key = jax.random.PRNGKey(0)
noise = jax.random.gumbel(key, logits.shape)
index = jnp.argmax(logits + noise)
result = index.astype(jnp.float32)
# A scalar containing one of 0, 1 or 2, not a vector of probabilities.

A class with a smaller logit can win because of its noise draw. Over many fresh draws, larger logits win more often. Two different seeds can still select the same class.

Problem

Implement gumbel_argmax(seed, logits) with this exact sequence:

  1. Create jax.random.PRNGKey(int(seed)).
  2. Draw jax.random.gumbel(key, logits.shape) with default float32 noise.
  3. Return jnp.argmax(logits + noise).astype(jnp.float32).

seed is a nonnegative integer-valued scalar supplied as a number. logits is a nonempty finite float32 vector of shape (K,). Return a float32 scalar containing the selected zero-based index. If perturbed scores tie in floating-point arithmetic, choose the first maximum, matching jnp.argmax.

The fixtures check the specified seeded draw. Another algorithm can sample from the same mathematical distribution without producing the same result for a particular seed.

Common pitfalls

  • Generate one noise value per logit. A single shared noise value leaves their order unchanged and cannot implement categorical sampling.
  • The same key, shape and random configuration reproduce the noise. Split keys for fresh draws; changing a seed does not guarantee a different selected class.
  • Return a scalar, not a length-one array or a one-hot vector. The float32 cast is this exercise's output convention; indices are usually integers.
  • Default float32 Gumbel noise has limited precision in its tails, so extremely unlikely events can be biased. The distributional identity is mathematical, not a claim of exact floating-point probabilities.
  • Python int(seed) requires a concrete seed. Mark the seed static when wrapping this exercise in jax.jit, or accept a JAX key directly in a reusable sampler.