fix(ops): propagate NaN through median

[devteamaegis]

Proposed changes

What's broken. mx.median silently drops NaN — it returns a real number for input that contains NaN.

import mlx.core as mx
mx.median(mx.array([1.0, float("nan"), 0.0]))   # array(1, dtype=float32)

Reproduced on CPU (mlx 0.32.1.dev20260810+e78d894, source build, -DMLX_BUILD_METAL=OFF). Both NumPy and PyTorch return nan here:

input	mx.median	torch.median	np.median
[1.0, nan, 0.0]	1.0	nan	nan
[nan, 1.0, 0.0]	1.0	nan	nan
[-5.0, nan, 3.0]	3.0	nan	nan
[1.0, 0.0, nan]	1.0	nan	nan

It is also inconsistent inside MLX: max, min, mean, cummax and cummin all propagate NaN, so median is the odd one out among the reductions.

The behaviour is shape-dependent, which makes it easy to miss. median over an axis of even length can average the NaN in by accident, so the same array gives a NaN along one axis and a plausible-looking number along another:

x = mx.array([[1.0, float("nan"), 3.0], [4.0, 5.0, 6.0]])
mx.median(x, axis=0)   # array([2.5, nan, 4.5])  <- correct, by luck
mx.median(x, axis=1)   # array([3, 5])           <- NaN dropped; numpy gives [nan, 5]
mx.median(x)           # array(4.5)              <- NaN dropped; numpy gives nan

Why. median sorts the reduced axes and slices the midpoint (mlx/ops.cpp). sort moves NaN to the end of the axis, so for an odd-length axis the midpoint is always a non-NaN element and the NaN is never observed.

The fix. After taking the midpoint, mask the result where the reduced axes contain a NaN. Guarded on issubdtype(a.dtype(), inexact), so integer input (which is promoted to float but can never be NaN) keeps the original code path. Complex is covered too, matching NumPy's (nan+0j).

The test. test_median_nan in python/tests/test_ops.py, covering odd/even axis lengths, NaN in leading/middle/trailing position, float16/bfloat16/float32, per-axis and all-axes reductions, keepdims, complex, and negative controls (NaN-free float input and integer input are unchanged).

Fails before, passes after:

# before
FAILED python/tests/test_ops.py::TestOps::test_median_nan - AssertionError: False is not true : [1.0, nan, 0.0] mlx.core.float16
1 failed, 1 passed, 149 deselected

# after
2 passed, 149 deselected

python/tests/test_ops.py, test_autograd.py and test_reduce.py are green (215 passed, 5422 subtests). The rest of the suite is green apart from 18 pre-existing Metal DLPack import is not available failures from my CPU-only build, which this change does not touch.

Benchmark. median already does an O(n log n) sort, so the added O(n) isnan + any pass is small. CPU, M4, best of 3 runs, using benchmarks/python/time_utils.py:

case	before	after	delta
median((1_000_000,)) all axes	77.12 ms	75.47 ms	−2.1%
median((1024,1024)) all axes	80.81 ms	79.55 ms	−1.6%
median((1024,1024), axis=1)	37.38 ms	38.69 ms	+3.5%
median((256,256,64), axis=(0,2))	224.56 ms	229.18 ms	+2.1%
median((64,64,64,64), axis=(1,3))	739.33 ms	785.11 ms	+6.2%
median((1024,1024)) int32, all axes	48.71 ms	50.07 ms	+2.8%

The int32 row takes the guarded path, so it runs byte-identical code in both builds; its +2.8% is the run-to-run noise floor on this machine. Everything except the largest 4-D multi-axis reduce sits inside that noise, and that case is +6.2%.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed) — no API change, and no other reduction documents its NaN behaviour, so the docstring is unchanged