From 9051982281b9c1de97f4c54234ee2121a9e31f11 Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Sat, 15 Aug 2026 12:43:56 -0700 Subject: [PATCH 1/3] Start complex min and max scans from an infinite identity An exclusive scan writes the identity of its operation into the first position. That identity came from std::numeric_limits, which has no specialization for complex64_t and so handed back zero. Zero then won every comparison against a negative real part, and cummax and cummin returned all zeros. The CUDA backend already uses an infinite pair. --- mlx/backend/cpu/scan.cpp | 27 ++++++++++++++++++++------- python/tests/test_ops.py | 16 ++++++++++++++++ 2 files changed, 36 insertions(+), 7 deletions(-) diff --git a/mlx/backend/cpu/scan.cpp b/mlx/backend/cpu/scan.cpp index 93e67825e2..92cd68ed49 100644 --- a/mlx/backend/cpu/scan.cpp +++ b/mlx/backend/cpu/scan.cpp @@ -191,6 +191,23 @@ void scan_op( } } +// The identity an exclusive min or max scan starts from. std::numeric_limits +// is not specialized for complex64_t, so complex used to start from zero, +// which then won every comparison against a negative real part. The CUDA +// backend already uses an infinite pair here. +template +U scan_extreme(const Dtype& dtype, bool maximum) { + constexpr auto inf = std::numeric_limits::infinity(); + if constexpr (std::is_same_v) { + return maximum ? complex64_t{inf, inf} : complex64_t{-inf, -inf}; + } else if (issubdtype(dtype, floating)) { + return maximum ? static_cast(inf) : static_cast(-inf); + } else { + return maximum ? std::numeric_limits::max() + : std::numeric_limits::min(); + } +} + template void scan_dispatch( Scan::ReduceType rtype, @@ -221,9 +238,7 @@ void scan_dispatch( } return x < y ? x : y; }; - auto init = (issubdtype(in.dtype(), floating)) - ? static_cast(std::numeric_limits::infinity()) - : std::numeric_limits::max(); + auto init = scan_extreme(in.dtype(), /* maximum = */ true); scan_op(in, out, axis, reverse, inclusive, op, init); break; } @@ -236,9 +251,7 @@ void scan_dispatch( } return x < y ? y : x; }; - auto init = (issubdtype(in.dtype(), floating)) - ? static_cast(-std::numeric_limits::infinity()) - : std::numeric_limits::min(); + auto init = scan_extreme(in.dtype(), /* maximum = */ false); scan_op(in, out, axis, reverse, inclusive, op, init); break; } @@ -246,7 +259,7 @@ void scan_dispatch( auto op = [](U a, T b) { return detail::LogAddExp{}(a, static_cast(b)); }; - auto init = (issubdtype(in.dtype(), floating)) + auto init = (issubdtype(in.dtype(), inexact)) ? static_cast(-std::numeric_limits::infinity()) : std::numeric_limits::min(); scan_op(in, out, axis, reverse, inclusive, op, init); diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 22a794b4fc..98cd606d65 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2677,6 +2677,22 @@ def test_scan_size_one_axis(self): out = getattr(mx, op)(a, axis=0) self.assertTrue(np.array_equal(np.array(out), expected)) + def test_scans_complex_exclusive(self): + # An exclusive scan starts from the identity of its operation, and + # numeric_limits has none for complex, so it used to start from zero + # and swallow every negative real part. + a = mx.array([-3 + 1j, -1 + 2j, -4 + 0j, 0 + 5j, 2 - 1j]) + for op in ("cummax", "cummin", "logcumsumexp"): + mxop = getattr(mx, op) + for reverse in (False, True): + inclusive = mxop(a, axis=0, inclusive=True, reverse=reverse) + exclusive = mxop(a, axis=0, inclusive=False, reverse=reverse) + if reverse: + got, want = exclusive[:-1], inclusive[1:] + else: + got, want = exclusive[1:], inclusive[:-1] + self.assertTrue(mx.allclose(got, want), msg=f"{op} reverse={reverse}") + def test_cummax_cummin_nan(self): nan = float("nan") cases = [ From 7cdd31f2830b7b53b54694b7ec520d09440d5b23 Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Sat, 15 Aug 2026 14:40:33 -0700 Subject: [PATCH 2/3] Note that Metal also specializes the complex limits --- mlx/backend/cpu/scan.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx/backend/cpu/scan.cpp b/mlx/backend/cpu/scan.cpp index 92cd68ed49..a698060c7e 100644 --- a/mlx/backend/cpu/scan.cpp +++ b/mlx/backend/cpu/scan.cpp @@ -193,8 +193,8 @@ void scan_op( // The identity an exclusive min or max scan starts from. std::numeric_limits // is not specialized for complex64_t, so complex used to start from zero, -// which then won every comparison against a negative real part. The CUDA -// backend already uses an infinite pair here. +// which then won every comparison against a negative real part. The Metal and +// CUDA backends already use an infinite pair here. template U scan_extreme(const Dtype& dtype, bool maximum) { constexpr auto inf = std::numeric_limits::infinity(); From 693ff89bdbc8babc47ef5cb0149545dd6f920ef2 Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Mon, 17 Aug 2026 11:01:11 -0700 Subject: [PATCH 3/3] Pin the cpu stream in the complex scan test The identity being fixed lives in backend/cpu/scan.cpp, but the default stream is the gpu, which has its own identity and was already correct. Without pinning, the test passed on macOS without exercising the fix. Both streams are checked now. --- python/tests/test_ops.py | 35 ++++++++++++++++++++++++----------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 98cd606d65..78d8c872ff 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2680,18 +2680,31 @@ def test_scan_size_one_axis(self): def test_scans_complex_exclusive(self): # An exclusive scan starts from the identity of its operation, and # numeric_limits has none for complex, so it used to start from zero - # and swallow every negative real part. + # and swallow every negative real part. Only the exclusive form reads + # the identity at all; the inclusive one starts from element 0. + # + # The cpu stream is pinned because the identity being fixed lives in + # mlx/backend/cpu/scan.cpp and the default stream is the gpu, which has + # its own. Both are checked so neither backend can regress. a = mx.array([-3 + 1j, -1 + 2j, -4 + 0j, 0 + 5j, 2 - 1j]) - for op in ("cummax", "cummin", "logcumsumexp"): - mxop = getattr(mx, op) - for reverse in (False, True): - inclusive = mxop(a, axis=0, inclusive=True, reverse=reverse) - exclusive = mxop(a, axis=0, inclusive=False, reverse=reverse) - if reverse: - got, want = exclusive[:-1], inclusive[1:] - else: - got, want = exclusive[1:], inclusive[:-1] - self.assertTrue(mx.allclose(got, want), msg=f"{op} reverse={reverse}") + devices = [mx.cpu] + if mx.default_device() != mx.cpu: + devices.append(mx.default_device()) + for device in devices: + with mx.stream(device): + for op in ("cummax", "cummin", "logcumsumexp"): + mxop = getattr(mx, op) + for reverse in (False, True): + inclusive = mxop(a, axis=0, inclusive=True, reverse=reverse) + exclusive = mxop(a, axis=0, inclusive=False, reverse=reverse) + if reverse: + got, want = exclusive[:-1], inclusive[1:] + else: + got, want = exclusive[1:], inclusive[:-1] + self.assertTrue( + mx.allclose(got, want), + msg=f"{op} reverse={reverse} device={device}", + ) def test_cummax_cummin_nan(self): nan = float("nan")