zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix max complex64 nan propagation and add test

Browse Source
This commit is contained in:
Joona Havukainen
2025-07-06 15:26:09 -07:00
parent 5b089dc5da
commit e0bb9f3ef8
2 changed files with 44 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
mlx/backend/metal/kernels/reduction/ops.h
View File

@@ -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<complex64_t>(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<float>(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<float>(NAN));
} else {
return complex64_t(static_cast<float>(NAN), static_cast<float>(NAN));
}
return a > b ? a : b;
}
};

31
python/tests/test_reduce.py
View File

@@ -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)