diff --git a/mlx/ops.cpp b/mlx/ops.cpp index b8973e8ad8..066e288260 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2324,6 +2324,16 @@ array median( array(0.5, dtype), s); } + // Sorting moves NaN to the end, so the midpoint slice never selects it. + // Propagate it explicitly to stay consistent with max, min and mean. + if (issubdtype(a.dtype(), inexact)) { + median_a = where( + any(isnan(flat_a, s), -1, /* keepdims = */ true, s), + array(std::numeric_limits::quiet_NaN(), dtype), + median_a, + s); + } + median_a = squeeze(median_a, -1, s); if (keepdims) { median_a = expand_dims(median_a, sorted_axes, s); diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 14dab531cf..5bc4e1ef60 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -962,6 +962,38 @@ def test_median(self): out_np = np.median(x, axis=(0, 1, 3), keepdims=True) self.assertTrue(np.allclose(out, out_np)) + def test_median_nan(self): + nan = float("nan") + + # Odd and even lengths, with the NaN in a few different positions. + for vals in ([1.0, nan, 0.0], [nan, 1.0, 0.0], [1.0, 2.0, nan, 4.0]): + for dtype in (mx.float16, mx.bfloat16, mx.float32): + out = mx.median(mx.array(vals, dtype=dtype)) + self.assertTrue(mx.isnan(out).item(), msg=f"{vals} {dtype}") + + x = mx.array([[1.0, nan, 3.0], [4.0, 5.0, 6.0]]) + self.assertTrue( + np.array_equal( + np.array(mx.median(x, axis=1)), np.median(x, axis=1), equal_nan=True + ) + ) + self.assertTrue( + np.array_equal( + np.array(mx.median(x, axis=0)), np.median(x, axis=0), equal_nan=True + ) + ) + self.assertTrue(mx.isnan(mx.median(x)).item()) + self.assertEqual(mx.median(x, axis=1, keepdims=True).shape, (2, 1)) + + # Complex NaN propagates too, matching NumPy. + out = mx.median(mx.array([complex(1, 0), complex(nan, 0), complex(0, 0)])) + self.assertTrue(mx.isnan(out).item()) + + # A NaN-free array is unaffected, and integers are never NaN. + x = mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + self.assertTrue(np.allclose(mx.median(x, axis=1), np.median(x, axis=1))) + self.assertEqual(mx.median(mx.array([0, 1, 2, 3, 4])).item(), 2) + def test_var(self): x = mx.array( [