Fix BatchNorm running variance estimator

[ishtihoss]

Summary

• update BatchNorm running_var with the unbiased variance estimate while keeping training normalization on the existing biased batch variance
• preserve finite running stats for the one-value-per-channel case where unbiased variance is undefined
• cover 2D, 3D, 4D, and degenerate BatchNorm running-stat behavior

Tests

• black --check python/mlx/nn/layers/normalization.py python/tests/test_nn.py
• python3 -m py_compile python/mlx/nn/layers/normalization.py python/tests/test_nn.py
• git diff --check upstream/main..HEAD
• PYTHONPATH=python:python/tests /Users/porkr/projects/contributor/.venv/bin/python -m unittest python.tests.test_nn.TestLayers.test_batch_norm python.tests.test_nn.TestLayers.test_batch_norm_stats -v

[Pablosinyores]

I looked into this since I've been in these normalization layers recently, and I think the PR is actually consistent with the paper — the unbiased estimate is applied only to running_var, not to the training-time normalization.

Walking __call__:

• mean, var = self._calc_stats(x) gives the biased batch variance (ddof=0).
• In the self.training branch the PR derives running_var = var * n/(n-1) (unbiased) and folds only that into self.running_var.
• The normalization a few lines down, x = (x - mean) * mx.rsqrt(var + self.eps), still uses the biased var — that line is untouched.

So training normalization stays on the biased batch variance and only the inference statistic becomes unbiased, which matches Ioffe & Szegedy (inference Var[x] = m/(m-1)·E[σ²_B]) and PyTorch, whose docs state the batch is normalized with the biased estimator while running_var stores the unbiased one.

Quick sanity check (n=5): var(x, ddof=1)/var(x, ddof=0) is exactly 1.25 = n/(n-1), so the stats_size/(stats_size-1) factor is the standard Bessel correction, and the stats_size == 1 guard keeps running_var finite where the unbiased estimate is undefined.

Might be worth a one-line code comment on why running_var uses the corrected estimate while the batch is normalized with the biased one — that asymmetry is exactly the part that's easy to misread.

[zcbenz]

Thanks a lot for looking into this! After checking PyTorch's implementation I agree this PR does the right fix.

On the implementation I think we can use mx.var(ddof=1) to get unbiased variance directly. For the tests I do prefer verification against PyTorch's implementation rather than reimplementing the algorithm in the test code.

[ishtihoss]

Thanks, updated in 9a4c1be:\n\n- switched the running_var update to guarded mx.var(..., ddof=1) while keeping training normalization on the biased batch variance\n- added a PyTorch parity test for 2D, 3D, and 4D BatchNorm covering training outputs, running stats, and eval outputs\n- kept the one-value-per-channel fallback as an MLX-only test since PyTorch raises for that case

[zcbenz]

Thanks!