Skip to content

← JAX Stochasticity step 7 of 25

Medium Primitives

Dirichlet Sampling

Jump to the task

Draw a whole probability vector

A Dirichlet draw divides a total of one among K categories. For example, [0.2, 0.5, 0.3] is one possible three-category probability vector. A categorical draw would instead select a single category index. Dirichlet samples describe distributions, not class labels.

The coordinates of a row are related: if one takes more of the total, less remains for the others. Separate rows are separate draws. The set of nonnegative K-component vectors that sum to one is the probability simplex. Its dimension is K - 1; for three categories, it is a triangle, not a three-dimensional simplex.

Alpha sets both the mean and the spread

Alpha contains positive concentration parameters. They need not sum to one. Let total = sum(alpha); the mean of coordinate i is alpha[i] / total. Multiplying all entries by the same positive factor keeps that mean but changes the spread.

  • [1, 1, 1] gives uniform density across the triangle of probability vectors. It does not return [1/3, 1/3, 1/3] every time.
  • [20, 20, 20] concentrates draws more tightly around equal proportions.
  • [0.1, 0.1, 0.1] favors vectors near corners, often with one dominant entry. The entries are not required to be exactly zero or one.
  • [50, 1, 1] has mean [50/52, 1/52, 1/52]. The first coordinate's expected share is about 96%, not a fixed share in every row.

Large asymmetric concentrations gather around their normalized mean, not necessarily the center. Do not normalize alpha before sampling: doing so preserves the mean but changes the distribution's spread.

A concrete draw

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
alpha = jnp.array([1., 1., 1.])
samples = jax.random.dirichlet(key, alpha, shape=(2,))
# approximately [[0.364721, 0.428289, 0.206990],
#                [0.494387, 0.325759, 0.179855]]

This draws two three-component vectors, so the output shape is (2, 3). Each row sums to one up to floating-point rounding. Neither each column nor the entire matrix must sum to one. The example uses float32 and the runtime's default Threefry and Dirichlet sampling settings.

In Bayesian categorical modeling, adding observed category counts to alpha gives the posterior Dirichlet parameters. The distribution is also useful when an application needs random mixture weights.

Problem

Implement dirichlet_sample(seed, alpha, n) by creating jax.random.PRNGKey(int(seed)) and calling jax.random.dirichlet once with shape=(int(n),).

Seed is a whole Python number in [0, 2**32), supplied as a float. Alpha is a nonempty one-dimensional float32 array of finite positive concentrations. N is a nonnegative whole Python number supplied as a float. Return a float32 array (int(n), len(alpha)). Zero draws returns shape (0, len(alpha)); with one category, every row contains just 1.

Use the original key and default runtime sampler. Do not split the key, normalize alpha, or replace the draw with a repeated mean vector. For JIT, mark seed and n static: jax.jit(dirichlet_sample, static_argnames=("seed", "n")). Alpha values can remain dynamic; int() does not make traced inputs static.