diff --git a/mlx/backend/cpu/simd/base_simd.h b/mlx/backend/cpu/simd/base_simd.h index 41193890aa..626dd93df9 100644 --- a/mlx/backend/cpu/simd/base_simd.h +++ b/mlx/backend/cpu/simd/base_simd.h @@ -244,7 +244,7 @@ Simd maximum(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; if constexpr (!std::is_integral_v) { - if (std::isnan(a)) { + if (mlx::core::isnan(a)) { return a; } } @@ -256,7 +256,7 @@ Simd minimum(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; if constexpr (!std::is_integral_v) { - if (std::isnan(a)) { + if (mlx::core::isnan(a)) { return a; } } diff --git a/mlx/backend/cpu/sort.cpp b/mlx/backend/cpu/sort.cpp index cf8d86de34..92378981b6 100644 --- a/mlx/backend/cpu/sort.cpp +++ b/mlx/backend/cpu/sort.cpp @@ -2,7 +2,6 @@ #include #include -#include #include #include "mlx/backend/common/utils.h" @@ -19,9 +18,9 @@ namespace { template bool nan_aware_less(T a, T b) { if constexpr (is_floating_point_v || std::is_same_v) { - if (std::isnan(a)) + if (mlx::core::isnan(a)) return false; - if (std::isnan(b)) + if (mlx::core::isnan(b)) return true; } return a < b; @@ -206,10 +205,10 @@ void argsort(const array& in, array& out, int axis) { auto v2 = data_ptr[b * in_stride]; // Handle NaNs (place them at the end) - if constexpr (is_floating_point_v) { - if (std::isnan(v1)) + if constexpr (is_floating_point_v || std::is_same_v) { + if (mlx::core::isnan(v1)) return false; - if (std::isnan(v2)) + if (mlx::core::isnan(v2)) return true; } @@ -315,10 +314,10 @@ void argpartition(const array& in, array& out, int axis, int kth) { auto v2 = data_ptr[b * in_stride]; // Handle NaNs (place them at the end) - if constexpr (is_floating_point_v) { - if (std::isnan(v1)) + if constexpr (is_floating_point_v || std::is_same_v) { + if (mlx::core::isnan(v1)) return false; - if (std::isnan(v2)) + if (mlx::core::isnan(v2)) return true; } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index d5568610e8..a89bf7f8b0 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -593,6 +593,11 @@ def test_minimum(self): self.assertTrue(math.isnan(mx.minimum(a, b).item())) self.assertTrue(math.isnan(mx.minimum(b, a).item())) + # A NaN in either part of the first argument propagates. + a = mx.array([complex(3, float("nan"))]) + b = mx.array([2 + 0j]) + self.assertTrue(mx.array_equal(mx.minimum(a, b), a, equal_nan=True)) + def test_maximum(self): x = mx.array([0.0, -5, 10.0]) y = mx.array([1.0, -7.0, 3.0]) @@ -605,6 +610,11 @@ def test_maximum(self): self.assertTrue(math.isnan(mx.maximum(a, b).item())) self.assertTrue(math.isnan(mx.maximum(b, a).item())) + # A NaN in either part of the first argument propagates. + a = mx.array([complex(1, float("nan"))]) + b = mx.array([2 + 0j]) + self.assertTrue(mx.array_equal(mx.maximum(a, b), a, equal_nan=True)) + def test_floor(self): x = mx.array([-22.03, 19.98, -27, 9, 0.0, -np.inf, np.inf]) expected = [-23, 19, -27, 9, 0, -np.inf, np.inf] @@ -4274,6 +4284,11 @@ def test_sort_nan(self): expected = mx.array([0.0 + 1j, 2.0 + 1j, 3.0 + 1j, mx.nan + 2j]) self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True)) + # A NaN in the imaginary part also sorts last. + x = mx.array([3.0 + 1j, complex(1.0, mx.nan), 2.0 + 1j, 0.0 + 1j]) + expected = mx.array([0.0 + 1j, 2.0 + 1j, 3.0 + 1j, complex(1.0, mx.nan)]) + self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True)) + def test_argsort_nan(self): for dtype in [mx.float32, mx.float16, mx.bfloat16]: with self.subTest(dtype=dtype): @@ -4283,6 +4298,13 @@ def test_argsort_nan(self): sorted_x = mx.take(x, indices) self.assertTrue(mx.array_equal(sorted_x, expected, equal_nan=True)) + # Also test complex values + for nan_val in [mx.nan + 2j, complex(1.0, mx.nan)]: + x = mx.array([3.0 + 1j, nan_val, 2.0 + 1j, 0.0 + 1j]) + expected = mx.array([0.0 + 1j, 2.0 + 1j, 3.0 + 1j, nan_val]) + sorted_x = mx.take(x, mx.argsort(x)) + self.assertTrue(mx.array_equal(sorted_x, expected, equal_nan=True)) + def test_to_from_fp8(self): vals = mx.array( [448, 256, 192, 128, 96, 64, 48, 32, 24, 16, 12, 8, 6, 4, 3, 2, 0.015625]