Skip to content

← JAX Stochasticity step 22 of 25

Medium Primitives

Multivariate Normal Sampling

Jump to the task

Draw related quantities together

A multivariate normal draw is a vector. Its mean specifies the center of each coordinate, and its covariance describes how the coordinates vary together. Separate rows are separate draws from that distribution; coordinates within one row can be correlated.

The diagonal of cov contains variances, not standard deviations. For cov = [[4, 1], [1, 4]], both coordinates have standard deviation 2. Their correlation is 1 / (2 * 2) = 0.25. Off-diagonal covariance entries equal correlations only when both standard deviations are one.

Ask for rows of vectors

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
mean = jnp.zeros(2)
cov = jnp.eye(2)
samples = jax.random.multivariate_normal(key, mean, cov, shape=(3,))
# approximately [[ 1.622642,  2.025265],
#                [-0.433594, -0.078617],
#                [ 0.176091, -0.972089]]

The final dimension comes from mean.shape. Here shape=(3,) requests three draws of a two-dimensional vector, giving (3, 2). Passing shape=(3, 2) would add an unwanted dimension and give (3, 2, 2). This example uses the default Threefry configuration and float32.

The default factorization needs positive definiteness

The default method is Cholesky. For a symmetric positive-definite covariance, it finds a lower-triangular L such that cov = L @ L.T. One column-vector draw can then be written mean + L @ z, where z contains independent standard-normal values. For draws stored as rows, the corresponding expression is mean + Z @ L.T.

Symmetry and positive diagonal entries alone are not enough. For example, [[1, 2], [2, 1]] has those properties but has a negative eigenvalue, so it is not a valid covariance. Invalid inputs can give NaNs rather than raising an exception.

A positive-semidefinite covariance can be singular and still define a Gaussian distribution, but it may fail with the default Cholesky method. JAX offers eigh and svd methods for singular covariance cases. This exercise supplies positive-definite matrices and uses Cholesky; changing factorization can change the samples for a fixed key.

Problem

Implement multivariate_normal_sample(seed, mean, cov, n) using jax.random.PRNGKey(int(seed)) and one call to jax.random.multivariate_normal with shape=(int(n),).

  • Seed is a whole Python number in [0, 2**32), supplied as a float.
  • mean is a finite float32 vector of shape (d,), with d >= 1.
  • cov is a finite, symmetric positive-definite float32 matrix (d, d).
  • n is a nonnegative whole Python number supplied as a float.
  • Return a float32 array (int(n), d). Zero draws returns shape (0, d).

Use the default runtime sampler, Cholesky method and seed-derived key without adding a split. The sample mean and covariance need not exactly match their parameters for a small number of rows. Do not normalize the result to force them to match.

For JIT, mark seed and n static: jax.jit(multivariate_normal_sample, static_argnames=("seed", "n")). Mean and covariance values can remain dynamic. int() cannot make a traced seed or output length static.