Skip to content

← JAX Stochasticity step 5 of 25

Medium Primitives

Reparameterization Trick: Gaussian

Jump to the task

Why this matters

To sample a Gaussian with mean mu and standard deviation sigma, first draw standard normal noise eps, then compute mu + sigma * eps. For fixed noise, changing either parameter changes the sample smoothly. This lets autodiff follow how a loss changes along that sample path.

Gaussian latent-variable models, including many VAEs, use this construction to estimate gradients. Here the scale parameter is log_sigma, the natural logarithm of the standard deviation. It is not the log variance.

Worked mini-example

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
mu, log_sigma = 2.0, 0.5
eps = jax.random.normal(key, (3,))
sigma = jnp.exp(log_sigma)
z = mu + sigma * eps
# Mean 2, standard deviation exp(0.5) ~= 1.649.
# Variance exp(1) ~= 2.718, not 1.649.

Holding eps fixed gives dz/dmu = 1 and dz/dlog_sigma = sigma * eps, element by element. The noise distribution itself does not depend on either parameter.

Problem

Implement reparam_normal(seed, mu, log_sigma, n):

  1. Create jax.random.PRNGKey(int(seed)).
  2. Draw jax.random.normal(key, (int(n),)).
  3. Return mu + jnp.exp(log_sigma) * eps.

seed is a nonnegative integer-valued scalar. n is a nonnegative integer-valued sample count. Both may arrive as floating-point numbers. mu and log_sigma are finite scalars; assume the scale and resulting samples remain finite and the scale is nonzero in float32. With the default float32 random configuration, return an array of shape (int(n),). A zero sample count gives an empty array.

The tests compare the specified seeded samples, not just their approximate mean and variance. They do not by themselves check differentiability.

Common pitfalls

  • Standard deviation is exp(log_sigma); variance is exp(2 * log_sigma). If a different model supplies log variance, its conversion is different.
  • A Gaussian sampling API does not inherently block gradients. What matters is whether its implementation exposes a differentiable path through the parameters. This exercise makes that path explicit.
  • Reusing a key reproduces the noise for the same shape and configuration. Split keys when independent draws are needed.
  • PRNGKey can accept an integer scalar array; it does not require a Python int in general. The explicit int(seed) and int(n) conversions here do require concrete values. Mark seed and n static if using jax.jit.
  • Exponentiation enforces a positive scale in exact arithmetic, but extreme values can overflow or underflow in floating-point arithmetic.