diff --git a/mlx/backend/metal/kernels/reduction/ops.h b/mlx/backend/metal/kernels/reduction/ops.h index a3f3202fa..ee6fc5e5f 100644 --- a/mlx/backend/metal/kernels/reduction/ops.h +++ b/mlx/backend/metal/kernels/reduction/ops.h @@ -217,10 +217,19 @@ struct Max { template <> complex64_t operator()(complex64_t a, complex64_t b) { - if (metal::isnan(a.real) || metal::isnan(a.imag) || metal::isnan(b.real) || - metal::isnan(b.imag)) { - return static_cast(NAN); + bool real_is_nan = metal::isnan(a.real) || metal::isnan(b.real); + bool imag_is_nan = metal::isnan(a.imag) || metal::isnan(b.imag); + + if (!real_is_nan && !imag_is_nan) { + return a > b ? a : b; + } else if (real_is_nan && !imag_is_nan) { + return complex64_t( + static_cast(NAN), a.imag > b.imag ? a.imag : b.imag); + } else if (!real_is_nan && imag_is_nan) { + return complex64_t( + a.real > b.real ? a.real : b.real, static_cast(NAN)); + } else { + return complex64_t(static_cast(NAN), static_cast(NAN)); } - return a > b ? a : b; } }; diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py index 6b2ff06b8..b762aca50 100644 --- a/python/tests/test_reduce.py +++ b/python/tests/test_reduce.py @@ -179,6 +179,37 @@ class TestReduce(mlx_tests.MLXTestCase): ref = getattr(np, op)(x_np, axis=axis) self.assertTrue(np.array_equal(out, ref, equal_nan=True)) + def test_nanpropagation_complex64(self): + complex_array_1 = mx.array( + [1 + 1j, 2 + 2j, 3 + 3j, mx.nan + 4j], dtype=mx.complex64 + ).reshape(2, 2) + complex_array_2 = mx.array( + [1 + 1j, 2 + 2j, 3 + mx.nan * 1j, 4 + 4j], dtype=mx.complex64 + ).reshape(2, 2) + complex_array_3 = mx.array( + [1 + 1j, 2 + mx.nan * 1j, 3 + 3j, 4 + 4j], dtype=mx.complex64 + ).reshape(2, 2) + complex_array_4 = mx.array( + [mx.nan + 1j, 2 + 2j, 3 + 3j, 4 + 4j], dtype=mx.complex64 + ).reshape(2, 2) + + np_arrays = [ + np.array(complex_array_1), + np.array(complex_array_2), + np.array(complex_array_3), + np.array(complex_array_4), + ] + + for mx_arr, np_arr in zip( + [complex_array_1, complex_array_2, complex_array_3, complex_array_4], + np_arrays, + ): + for axis in [0, 1]: + for op in ["max"]: + out = getattr(mx, op)(mx_arr, axis=axis) + ref = getattr(np, op)(np_arr, axis=axis) + self.assertTrue(np.array_equal(out, ref, equal_nan=True)) + if __name__ == "__main__": mlx_tests.MLXTestRunner(failfast=True)