Fix BatchNorm running variance estimator - #3817

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.

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.

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

Hi @zcbenz, thanks again for the approval. I synced the branch with main, and the new CI run is waiting for maintainer approval. When convenient, could you approve the run and squash-merge once it's green?