Adding benchmarks and testing for max op nanpropagation

benchmarks/cpp/single_ops.cpp

@@ -192,6 +192,17 @@ void time_reductions() {
 
   auto argmin_along_1 = [&a]() { return mx::argmin(a, 1, false); };
   TIME(argmin_along_1);
+
+  auto indices = mlx::core::array({1});
+  auto updates = mlx::core::reshape(mlx::core::array({NAN}), {1, 1, 1});
+  std::vector<int> axes{0};
+  auto b = scatter(a, {indices}, updates, axes);
+  mx::eval(b);
+
+  auto max_along_0 = [&b]() { return mx::max(b, 0, false); };
+  TIME(max_along_0);
+  auto max_along_1 = [&b]() { return mx::max(b, 1, false); };
+  TIME(max_along_1);
 }
 
 void time_gather_scatter() {

benchmarks/python/single_ops.py

@@ -50,6 +50,11 @@ def time_maximum():
     mx.eval(a, b)
     time_fn(mx.maximum, a, b)
 
+def time_max():
+    a = mx.random.uniform(shape=(32, 1024, 1024))
+    a[1,1] = mx.nan
+    mx.eval(a)
+    time_fn(mx.max, a, 0)
 
 def time_negative():
     a = mx.random.uniform(shape=(10000, 1000))
@@ -108,6 +113,7 @@ if __name__ == "__main__":
 
     time_add()
     time_matmul()
+    time_max()
     time_maximum()
     time_exp()
     time_negative()

python/tests/test_reduce.py

@@ -153,6 +153,31 @@ class TestReduce(mlx_tests.MLXTestCase):
         x = x.transpose(1, 0, 2, 3, 4, 5, 6, 7, 8, 9)
         check(x, (1, 3, 5, 7, 9))
 
+    def test_nanpropagation(self):
+        dtypes = [
+            "uint8",
+            "uint16",
+            "uint32",
+            "int8",
+            "int16",
+            "int32",
+            "float16",
+            "float32",
+        ]
+
+        for dtype in dtypes:
+            with self.subTest(dtype=dtype):
+                x = (mx.random.normal((4, 4))).astype(getattr(mx, dtype))
+                indices = mx.random.randint(0, 4, shape=(6,)).reshape(3,2)
+                for idx in indices:
+                    x[*idx] = mx.nan
+                x_np = np.array(x)
+
+                for op in ["max"]:
+                    for axis in [0, 1]:
+                        out = getattr(mx, op)(x, axis=axis)
+                        ref = getattr(np, op)(x_np, axis=axis)
+                        self.assertTrue(np.array_equal(out, ref, equal_nan=True))
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner(failfast=True)

tests/ops_tests.cpp

@@ -1025,7 +1025,7 @@ TEST_CASE("test reduction ops") {
     CHECK(array_equal(min(x, 1), array({true, false})).item<bool>());
     CHECK(array_equal(min(x, 0), array({false, true, false})).item<bool>());
 
-    x = array({1.0f, NAN, 3.0f});
+    x = array({1.0f, NAN, 3.0f, 4.0f, 5.0f, 6.0f});
     CHECK(isnan(max(x).item<float>()));
 
   }