Start complex min and max scans from an infinite identity

mlx/backend/cpu/scan.cpp

@@ -190,6 +190,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 <typename U>
+U scan_extreme(const Dtype& dtype, bool maximum) {
+  constexpr auto inf = std::numeric_limits<float>::infinity();
+  if constexpr (std::is_same_v<U, complex64_t>) {
+    return maximum ? complex64_t{inf, inf} : complex64_t{-inf, -inf};
+  } else if (issubdtype(dtype, floating)) {
+    return maximum ? static_cast<U>(inf) : static_cast<U>(-inf);
+  } else {
+    return maximum ? std::numeric_limits<U>::max()
+                   : std::numeric_limits<U>::min();
+  }
+}
+
 template <typename T, typename U>
 void scan_dispatch(
     Scan::ReduceType rtype,
@@ -220,9 +237,7 @@ void scan_dispatch(
         }
         return x < y ? x : y;
       };
-      auto init = (issubdtype(in.dtype(), floating))
-          ? static_cast<U>(std::numeric_limits<float>::infinity())
-          : std::numeric_limits<U>::max();
+      auto init = scan_extreme<U>(in.dtype(), /* maximum = */ true);
       scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
       break;
     }
@@ -235,17 +250,15 @@ void scan_dispatch(
         }
         return x < y ? y : x;
       };
-      auto init = (issubdtype(in.dtype(), floating))
-          ? static_cast<U>(-std::numeric_limits<float>::infinity())
-          : std::numeric_limits<U>::min();
+      auto init = scan_extreme<U>(in.dtype(), /* maximum = */ false);
       scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
       break;
     }
     case Scan::LogAddExp: {
       auto op = [](U a, T b) {
         return detail::LogAddExp{}(a, static_cast<U>(b));
       };
-      auto init = (issubdtype(in.dtype(), floating))
+      auto init = (issubdtype(in.dtype(), inexact))
           ? static_cast<U>(-std::numeric_limits<float>::infinity())
           : std::numeric_limits<U>::min();
       scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);

python/tests/test_ops.py

@@ -2510,6 +2510,35 @@ def fn(its):
         mem4 = mx.get_peak_memory()
         self.assertEqual(mem2, mem4)
 
+    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 = [