From 919a26ef522454b6278031c029161e7b53c101d6 Mon Sep 17 00:00:00 2001
From: devteamaegis <devteam.aegis@gmail.com>
Date: Mon, 10 Aug 2026 12:26:51 -0400
Subject: [PATCH] fix(ops): propagate NaN through median

median sorts and takes the midpoint. Sorting moves NaN to the end of the
axis, so the midpoint slice never selects it and the NaN is silently
dropped. Mask the result on any(isnan(...)) over the reduced axes for
inexact dtypes, which matches max, min, mean, cummax and cummin, as well
as NumPy and PyTorch.
---
 mlx/ops.cpp              | 10 ++++++++++
 python/tests/test_ops.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

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<float>::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(
             [