Fix cpu exclusive scan for complex numbers

[ayaangazali]

An exclusive scan writes the identity of its operation into the first position. On the CPU that identity is picked like this:

auto init = (issubdtype(in.dtype(), floating))
    ? static_cast<U>(std::numeric_limits<float>::infinity())
    : std::numeric_limits<U>::max();

complex64 is inexact but not floating, so it takes the second branch. std::numeric_limits has no specialization for complex64_t, so max() returns a value initialized complex64_t, which is 0+0j. Complex comparison is lexicographic on the real part, so that zero beats every negative real part and the scan never recovers:

>>> a = mx.array([-3+1j, -1+2j, -4+0j, -5+5j])
>>> mx.cummax(a, axis=0, inclusive=True)
array([-3+1j, -1+2j, -1+2j, -1+2j], dtype=complex64)
>>> mx.cummax(a, axis=0, inclusive=False)
array([0+0j, 0+0j, 0+0j, 0+0j], dtype=complex64)

cummin goes the same way whenever the data is non negative, and logcumsumexp starts from 0+0j instead of -inf:

>>> b = mx.array([3+1j, 1+2j, 4+0j, 0+5j, 2-1j])
>>> mx.cummin(b, axis=0, inclusive=False)
array([0+0j, 0+0j, 0+0j, 0+0j, 0+0j], dtype=complex64)

The inclusive results are correct throughout, so the two disagree: an exclusive scan should equal the inclusive one shifted by a position, and for complex it does not.

Neither GPU backend has this problem: both Metal (mlx/backend/metal/kernels/utils.h) and CUDA (mlx/backend/cuda/device/utils.cuh) specialize Limits<complex_t<T>> to an infinite pair, so CumMax there already starts from {-inf, -inf} and CumMin from {inf, inf}. The CPU was the only backend without it. This gives the CPU the same identity. scan_extreme replaces the three copies of the ternary, and picks the pair {inf, inf} for complex so the identity also wins against an input whose real part is itself infinite, which a bare inf+0j would lose to on the imaginary tiebreak.

logcumsumexp keeps its own ternary, widened to inexact so complex reaches -inf. It deliberately does not use scan_extreme: its identity has to compare equal to the minval == -inf test inside LogAddExp, and an infinite imaginary part would fail that and produce NaN.

Verified that exclusive equals inclusive shifted by one across cumsum, cumprod, cummax, cummin and logcumsumexp, over 8 dtypes, shapes (7,), (3,5), (2,3,4), (1,) and (9,2), every axis, and both scan directions. Complex goes from broken to matching; the real dtypes are byte identical to before and still match numpy.

Out of scope and left alone: logcumsumexp on an unsigned integer input starts from 0 because there is no representable -inf, which is a different problem and predates this.

CPU only. test_ops.py, test_reduce.py, test_array.py, test_compile.py, test_vmap.py, test_autograd.py, test_nn.py and the C++ suite (249 cases, 3350 assertions) pass. The added test fails on main with cummax reverse=False.

I am a freshman working through this codebase and I used Claude Code alongside it. Every result above came from a run here.

[erwinzhang7]

Confirmed on an M5 Max, macOS 26.5, cpu stream, against c2bcf47:

cummax(z, axis=0, inclusive=False)[0]   0j  ->  -inf-infj
cummin(z, axis=0, inclusive=False)[0]   0j  ->   inf+infj

The inclusive scan is already correct on main: it starts
from element 0 and never reads the identity, so mx.cummax(z, axis=0) returns
the right answer without this patch. And the default stream is the GPU, so a
check has to pin mx.stream(mx.cpu) to reach mlx/backend/cpu/scan.cpp at all.
Only the exclusive form on the cpu stream exposes the 0+0j identity.

Test suites pass and it applies to current main on its own.

This is one of four open PRs fixing NaN handling in different ops. I applied all
four together and measured the result in
#4146 (comment)

[ayaangazali]

Thanks for running it, and the stream point was a real gap in my test rather than just a note.

You are right that the default stream is the gpu, and Metal already has the correct identity: utils.h specializes Limits<complex_t<T>> to an infinite pair, so the exclusive scan there was never wrong. That means my test was passing on macOS without touching the code it is supposed to guard. Fixed by pinning, and it now checks both:

devices = [mx.cpu]
if mx.default_device() != mx.cpu:
    devices.append(mx.default_device())
for device in devices:
    with mx.stream(device):
        ...

Confirmed it still catches the bug that way: reverting backend/cpu/scan.cpp to main makes it fail, where before the pin it would have passed on a gpu machine.

Your other observation matches what I found: only the exclusive form reads the identity, the inclusive one starts from element 0 and was already right on main. I put that in the comment so the next reader does not have to rediscover it.

[erwinzhang7]

Ran it again and all good!

[ayaangazali]

Rebased on main to clear the conflict with #4139 (both touched scan.cpp and test_ops.py). The two changes are independent: #4139 fixes the size-one axis row count in scan_op, this fixes the complex identity in scan_dispatch. Both new tests pass together locally.