Fix median dropping NaN

[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

[erwinzhang7]

Confirmed on an M5 Max, macOS 26.5, against c2bcf47. All four of your cases
return 1.0, 1.0, 3.0, 1.0 before, and nan after.

The consistency argument holds up when it is measured, and it is a bit tidier
than the description suggests. There are four open PRs fixing NaN handling in
different ops right now: this one, #4164 (from_fp8), #4272 (complex scan
identity) and #4291 (argmax/argmin). Applying all four on top of c2bcf47:

                 main    all four
max               nan       nan
min               nan       nan
mean              nan       nan
sum               nan       nan
var               nan       nan
median            1.0       nan     this PR
argmax -> value   1.0       nan     #4291
from_fp8(0x7f)  480.0       nan     #4164
cummax excl[0]     0j   -inf-infj   #4272

Five ops disagree with numpy on main; with the four applied, none do. So these
are not four unrelated fixes, they are the tail of one cleanup, and this is the
op with the widest blast radius of the four.

test_ops, test_reduce, test_array, test_autograd and test_fast pass with
all four applied, as does a 216-case argmax matrix over 12 dtypes, 9 shapes and
both streams. This PR also applies to current main on its own.

One caveat on what that does and does not show: I applied them in the order 4146,
4164, 4272, 4291, dropping #4291's merge-of-main commit since it was already
present. I did not test every merge order, so this says the four work together,
not that whichever lands last will rebase without help.