Implement BatchNorm with Mutable batch_stats

BatchNorm uses batch statistics during training and stored running statistics during evaluation. Implement both paths for a two-dimensional input, keeping the running statistics separate from the learned scale and bias.

Two collections of variables

For input shape (N, D), every feature has four scalar values:

• params/gamma: learned scale, initialized to one.
• params/beta: learned bias, initialized to zero.
• batch_stats/running_mean: stored mean, initialized to zero.
• batch_stats/running_var: stored variance, initialized to one.

Each array has shape (D,). Declare gamma and beta with self.param and the running statistics with self.variable. For example:

running_mean = self.variable(
    "batch_stats", "running_mean", lambda: jnp.zeros((d,))
)

Read or assign the array through running_mean.value. The collection name organizes the state; it does not automatically stop gradients. Training code normally differentiates with respect to params and updates running statistics separately. Other layers can also have non-parameter state.

Training uses the current batch

With use_running_average=False, compute each feature's statistics across axis=0:

mu = jnp.mean(x, axis=0)
var = jnp.mean((x - mu) ** 2, axis=0)
output = gamma * (x - mu) / jnp.sqrt(var + eps) + beta

Use population variance, dividing by N. With momentum 0.9, update the stored statistics after initialization:

new_running_mean = 0.9 * old_running_mean + 0.1 * mu
new_running_var  = 0.9 * old_running_var  + 0.1 * var

The training output uses the current batch statistics, not these moving averages. For one feature with batch values [1, 3], the batch mean is 2 and variance is 1. Starting from mean zero and variance one, the updated running values are 0.2 and 1.0.

Initialization and mutable application

Wrap the running-statistic assignments in if not self.is_initializing(). This preserves zero/one initial statistics instead of incorporating the initialization batch. Initialization permits state writes; omitting this guard would change the initial values, not inherently raise an error.

Later training calls must permit writes to the collection:

output, updates = model.apply(
    variables, x, use_running_average=False, mutable=["batch_stats"]
)
variables = {**variables, "batch_stats": updates["batch_stats"]}

The returned updates must be carried forward to retain the new statistics. The supplied variables mapping is not updated in place. Calling this training branch without mutable permission raises an error when it assigns the running values.

Evaluation uses stored statistics

With use_running_average=True, use the running mean and variance in the same normalization formula and leave them unchanged. Evaluation can run without mutable permission. Fixed running statistics make one example's result independent of the other examples in the evaluation batch.

This exercise initializes a fresh model on every wrapper call and discards returned updates. Its evaluation output therefore uses mean zero and variance one: x / sqrt(1 + 1e-5). It does not evaluate a previously trained model. In a training loop, retain the updated statistics as shown above.