Skip to content

← JAX Stochasticity step 2 of 25

Easy Primitives

Normal Sampling with mean and std

Jump to the task

Scale the spread, then shift the center

jax.random.normal draws standard-normal values: the distribution has mean zero and standard deviation one. To get a normal distribution with mean mu and standard deviation sigma, transform each draw z:

x = mu + sigma * z

Multiplying by sigma changes the spread; adding mu shifts the center. Sigma is the standard deviation, so the variance is sigma squared. For mu = 5 and sigma = 2, example z values [-1, 0, 1] become [3, 5, 7]. These illustrative values show the arithmetic, not a random draw.

Draw once, then transform

import jax

key = jax.random.PRNGKey(0)
z = jax.random.normal(key, (4,))
x = 5.0 + 2.0 * z
# approximately [8.2453, 9.0505, 4.1328, 4.8428]

This requests four values in one draw. It does not call the sampler four times with the same key, which would repeat the same scalar. The example uses the runtime's default Threefry configuration and float32 draws. The same key, shape and configuration reproduce z.

The sample mean of four values need not equal mu, and their sample standard deviation need not equal sigma. Those parameters describe the distribution, not exact statistics of each small batch. Do not normalize the sampled vector to force those statistics.

Common mistakes

  • normal has no mu or sigma arguments. Apply the transformation to its result.
  • sigma * (z + mu) scales the mean too. Use mu + sigma * z.
  • Multiply by sigma, not sigma squared. The latter uses the variance as a scale.
  • Use shape (int(n),), including the comma, to request a vector.

Holding z fixed makes this expression differentiable in mu and sigma: each output's derivative with respect to mu is 1, and with respect to sigma it is z. This is the basic idea behind Gaussian reparameterization. It does not mean every random sampling operation supports gradients with respect to all of its parameters.

Problem

Implement normal_mu_sigma(seed, n, mu, sigma) by creating jax.random.PRNGKey(int(seed)), drawing one standard-normal vector of shape (int(n),), and returning mu + sigma * z. Do not split the key: that would change the reproducible values expected here.

Seed is a whole Python number in [0, 2**32); n is a nonnegative whole Python number. Both arrive as floats and need conversion to int. Mu and sigma are finite real scalars, with sigma >= 0. Return a float32 array of shape (int(n),) using the default runtime sampling configuration. Sigma zero returns mu at every position; n zero returns an empty array.

For JIT, mark seed and n static: jax.jit(normal_mu_sigma, static_argnames=("seed", "n")). Mu and sigma can remain runtime values. int() does not make a traced seed or sample count static.