Fix cpu exclusive scan for complex numbers

mlx/backend/cpu/scan.cpp

@@ -191,6 +191,19 @@ void scan_op(
   }
 }
 
+template <typename U>
+U scan_init(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,
@@ -221,9 +234,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_init<U>(in.dtype(), /* maximum = */ true);
       scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
       break;
     }
@@ -236,17 +247,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_init<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

@@ -2691,6 +2691,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):
+        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 = [