Fix var/std for complex numbers

[ayaangazali]

mx.var on a complex input returns a complex number, and it can be negative. A variance cannot be either.

>>> x = mx.array([1+2j, -3-4j, 0.5-0.25j, 2+0j, -1+1j, 0+3j])
>>> mx.var(x)
array(-2.39062+4.34028j, dtype=complex64)
>>> mx.std(x)
array(1.13236+1.91647j, dtype=complex64)

numpy and torch both give 7.4600697 as float32 for the same input.

The cause is in var in mlx/ops.cpp:

auto v = sum(square(subtract(a, mu, s), s), axes, keepdims, s);

square computes z * z, but the variance of complex values is the mean squared magnitude, |z - mu|^2. For real inputs those are the same thing, so the bug is invisible until the input is complex, at which point the answer is silently wrong rather than an error. std inherits it since it is sqrt(var(...)).

This takes the magnitude of the deviations before squaring when the input is complex, and drops the accumulator to float32 so the output dtype matches numpy and torch. Real inputs take the same path as before, so there is no extra work on them.

abs is already instantiated for complex64 on Metal (instantiate_unary_base_same(Abs, complex64, complex64_t)) and CUDA, so this only composes ops that already run on every backend.

The gradient is unchanged: square's vjp 2|d| times abs's vjp sign(d) is 2d, exactly what square alone gave before.

Verified against numpy over shapes (6,), (3,4) and (2,3,4), every axis, ddof 0 and 1, and both keepdims settings, plus compile, vmap and grad. Empty inputs still give nan, and ddof >= N still gives nan, both unchanged. Added tests fail on main with mlx.core.complex64 != mlx.core.float32.

python/tests/test_ops.py, test_nn.py, test_reduce.py, test_autograd.py, test_compile.py, test_vmap.py, test_array.py and the C++ suite (249 cases, 3350 assertions) all pass. CPU-only build here, so I could not exercise the Metal or CUDA paths locally.

I am a freshman learning this codebase, and I used Claude Code while working through it. Every claim above I ran myself.

[zcbenz]

This is a bit different from JAX's implementation https://github.com/jax-ml/jax/blob/ac1ff2b9c112cbb41c59eca34a0d26eb6de410b8/jax/_src/numpy/reductions.py#L1145

which would be better?

[ayaangazali]

Good question, I went and checked both.

The JAX form is real(d * conj(d)), and mlx can't currently use it in an autodiff path. mx.grad through it raises [Primitive::vjp] Not implemented for Conjugate, so var and std on complex inputs would stop being differentiable. square(abs(d)) differentiates fine, gives 2d, and returns a clean 0 when a sample sits exactly on the mean instead of a NaN.

I also measured the other two axes on 0.29.3 so I wasn't trading accuracy or speed away for that. Relative error against a float64 numpy reference, 10k complex64 samples:

scale	square(abs)	real(d*conj)
1e-3	1.21e-08	8.28e-08
1e0	2.98e-08	2.98e-08
1e3	4.44e-08	5.84e-08
1e6	1.41e-08	1.41e-08

And elementwise over 4M values: square(abs) 2.21 ms, real(d*conj) 2.62 ms.

So abs-then-square is equal or better at every scale I tried and slightly faster, on top of being the only one of the two that survives grad. If you would rather match JAX's formulation I am happy to switch, it would just want a VJP for Conjugate landed first, otherwise this silently drops complex var/std out of autodiff.

[zcbenz]

Thanks for checking!

[ayaangazali]

I need to correct my answer above. I re-ran all of it against current main and most of what I told you was wrong.

Autodiff. My main claim, that real(d * conj(d)) would drop complex var/std out of autodiff, is false. Conjugate has had a VJP and JVP since #3386, and both forms differentiate and give the identical gradient 2d:

>>> a = mx.array([1+2j, -3-4j])
>>> mx.grad(lambda x: mx.real(x * mx.conj(x)).sum())(a)
array([2+4j, -6-8j], dtype=complex64)
>>> mx.grad(lambda x: mx.square(mx.abs(x)).sum())(a)
array([2+4j, -6-8j], dtype=complex64)

I got that error on an old 0.29.3 I had installed for the benchmark rather than on a build of main, and I did not re-check it before reporting it. That is on me.

Speed. Also wrong. End to end var on 4M complex64 is a wash, and my earlier elementwise number had the sign backwards:

full var, 4M   abs 7.95 ms    conj 7.79 ms
full var, 64k  abs 0.130 ms   conj 0.135 ms

Accuracy. My table measured var over 10k samples, so it was dominated by summation order rather than by the two formulas. Isolating the elementwise |z|^2 over 200k values, real(d * conj(d)) is in fact slightly the better of the two:

                 max rel err   mean rel err
square(abs)        1.754e-07      4.671e-08
real(d*conj)       1.169e-07      2.799e-08

Degenerate case. I claimed abs-then-square returns a clean 0 where the other gives NaN. Both return 0, with a 0 gradient, when every sample sits on the mean.

So there is no measured reason to prefer what this PR does over the JAX formulation. They are equivalent on gradients and speed, and the JAX one is marginally more accurate. Say the word and I will push the switch to real(d * conj(d)), or close this and let you take it whichever way you like. Sorry for the noise, I should have verified against main before answering.

[zcbenz]

No worries, let's just switch to JAX's algorithm in this branch.

[ayaangazali]

Switched to real(d * conj(d)), matching the JAX formulation you linked.

Re-verified after the change: complex var and std match numpy across shapes (6,), (3,4) and (2,3,4), every axis, ddof 0 and 1, and both keepdims settings, and still return float32. Real dtypes are untouched and keep their own dtype. The gradient is unchanged at 2d, and compile and vmap agree with eager. Empty input and ddof >= N still give nan.

test_ops.py, test_nn.py, test_reduce.py, test_autograd.py, test_compile.py, test_vmap.py, test_array.py and the C++ suite (249 cases, 3350 assertions) pass.