diff --git a/benchmarks/cpp/single_ops.cpp b/benchmarks/cpp/single_ops.cpp index 5b327be58..a5176e432 100644 --- a/benchmarks/cpp/single_ops.cpp +++ b/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 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() { diff --git a/benchmarks/python/single_ops.py b/benchmarks/python/single_ops.py index 3160a1833..24e609d92 100644 --- a/benchmarks/python/single_ops.py +++ b/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() diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py index 2b899c099..1617257d7 100644 --- a/python/tests/test_reduce.py +++ b/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) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index d3177a474..156d50838 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -1025,7 +1025,7 @@ TEST_CASE("test reduction ops") { CHECK(array_equal(min(x, 1), array({true, false})).item()); CHECK(array_equal(min(x, 0), array({false, true, false})).item()); - 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())); }