Skip to content

← Flax NNX step 24 of 100

Medium Primitives

NNX Implement LayerNorm

Jump to the task

Normalize each feature vector

LayerNorm centers and rescales a vector before applying a learned gain and offset. In this exercise, features occupy the last axis. For input shape (batch, tokens, D), each token has its own mean and variance over its D features. No statistics are shared between tokens or examples.

The layer has trainable state: gamma and beta, both shaped (D,). It has no running mean or variance. The same calculation is used during training and evaluation. BatchNorm can maintain running statistics and reduce across batch and spatial axes, so changing an axis alone does not reproduce all of BatchNorm's behavior.

The calculation

Compute mu = mean(x) and var = mean((x - mu) ** 2) along the last axis, retaining that axis with keepdims=True. This is population variance, divided by D, not the sample estimate divided by D - 1. Then return gamma * (x - mu) / sqrt(var + eps) + beta.

The positive eps keeps the denominator nonzero for a constant vector. Before the learned affine transform, the result has mean zero and variance var / (var + eps) in exact arithmetic, not exactly one. Learned per-feature gains and offsets can change both statistics.

A complete example

import jax.numpy as jnp
from flax import nnx

class MyLayerNorm(nnx.Module):
    def __init__(self, d, eps, rngs):
        self.gamma = nnx.Param(jnp.ones((d,), dtype=jnp.float32))
        self.beta = nnx.Param(jnp.zeros((d,), dtype=jnp.float32))
        self.eps = eps

    def __call__(self, x):
        mu = jnp.mean(x, axis=-1, keepdims=True)
        var = jnp.mean((x - mu) ** 2, axis=-1, keepdims=True)
        normalized = (x - mu) / jnp.sqrt(var + self.eps)
        return self.gamma[...] * normalized + self.beta[...]

model = MyLayerNorm(2, 3.0, nnx.Rngs(0))
x = jnp.array([[1., 3.], [5., 5.]], dtype=jnp.float32)
print(model(x))  # [[-0.5, 0.5], [0.0, 0.0]]
model.gamma[...] = jnp.array([2., 4.], dtype=jnp.float32)
model.beta[...] = jnp.array([1., -1.], dtype=jnp.float32)
print(model(x))  # [[0.0, 1.0], [1.0, -1.0]]

The deliberately large epsilon makes the arithmetic simple: the first row has mean 2 and variance 1, so its denominator is sqrt(1 + 3) = 2. The constant second row becomes zero before the affine transform, then returns beta. Epsilon is usually much smaller in model configurations.

rngs is accepted to match this exercise's constructor, but no random key is consumed. Ones and zeros do not require random initialization. keepdims=True leaves statistics shaped (..., 1) so they broadcast against the original input.

Loading visualization…

Problem

Implement layernorm_forward(seed, x, eps) using the module above:

  1. Build nnx.Rngs(int(seed)) and initialize a fresh module with d=x.shape[-1], gain ones, offset zeros, and the supplied epsilon.
  2. Compute mean and centered variance over the last axis.
  3. Return the normalized and affine-transformed array without flattening.

x is a nonempty float32 vector in the fixtures. The same implementation also supports leading batch or token dimensions. Use finite inputs whose squares and reductions stay within float32 range. eps is positive and finite. seed is a nonnegative integer-valued scalar, possibly supplied as a float; it has no effect on this deterministic initialization.

The output has the same shape as x. Each entrypoint call starts with fresh parameters and does no training. If JIT-compiling the entrypoint, make seed static because it is converted to a Python integer. Keep eps as a scalar in the array calculation rather than converting a traced value to Python float.

The numerical fixtures cannot prove that you used NNX parameters. Check that both gain and offset are nnx.Param objects and that a forward pass does not change them.