diff --git a/mlx/backend/cpu/scan.cpp b/mlx/backend/cpu/scan.cpp index 93e67825e2..a698060c7e 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 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(); + 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..78d8c872ff 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2677,6 +2677,35 @@ 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. 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]) + 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") cases = [