Fix the CPU scan over a size one axis with a padded stride

mlx/backend/cpu/scan.cpp

@@ -163,7 +163,8 @@ void scan_op(
     const Op& op,
     U init) {
   if (in.flags().row_contiguous) {
-    if (in.strides()[axis] == 1) {
+    // A size-one axis can carry any stride and still be row contiguous.
+    if (in.strides()[axis] == 1 || in.shape(axis) == 1) {
       contiguous_scan(
           in.data<T>(),
           out.data<U>(),

python/tests/test_ops.py

@@ -2631,6 +2631,20 @@ def fn(its):
         mem4 = mx.get_peak_memory()
         self.assertEqual(mem2, mem4)
 
+    def test_scan_size_one_axis(self):
+        # A size one axis can carry any stride and still be row contiguous, so
+        # the scan must not take its row count from that stride.
+        for op in ["cumsum", "cumprod", "cummax", "cummin"]:
+            for start in (1, 2, 3):
+                with self.subTest(op=op, start=start):
+                    base = mx.arange(1, 11, dtype=mx.float32).reshape(1, 10)
+                    a = base[:, start:]
+                    mx.eval(a)
+                    # The axis has size one, so an inclusive scan is the identity
+                    expected = np.array(a).copy()
+                    out = getattr(mx, op)(a, axis=0)
+                    self.assertTrue(np.array_equal(np.array(out), expected))
+
     def test_cummax_cummin_nan(self):
         nan = float("nan")
         cases = [