From 9b1a748047de3af04866e2516afeb51028510a28 Mon Sep 17 00:00:00 2001 From: Rohith Kapelli Date: Thu, 13 Aug 2026 14:35:18 +0530 Subject: [PATCH 1/2] Fix the CPU scan over a size one axis with a padded stride row_contiguous exempts size one axes, so such an axis can carry any stride and the array is not copied. The strided scan then takes its row count from that stride, size / shape[axis] / stride, which floors to zero and leaves the output unwritten. A padded stride is only reachable on a size one axis, where the scan is elementwise, so send that case to the contiguous path. --- mlx/backend/cpu/scan.cpp | 6 +++++- python/tests/test_ops.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cpu/scan.cpp b/mlx/backend/cpu/scan.cpp index 3ebbe0a3c3..23151c8e02 100644 --- a/mlx/backend/cpu/scan.cpp +++ b/mlx/backend/cpu/scan.cpp @@ -163,7 +163,11 @@ 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, so + // the strided count below would come out zero. The scan is elementwise + // in that case, which is what the contiguous path with a stride of one + // computes. + if (in.strides()[axis] == 1 || in.shape(axis) == 1) { contiguous_scan( in.data(), out.data(), diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index e5761c78b2..ccbc9390dc 100644 --- a/python/tests/test_ops.py +++ b/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, stream=mx.cpu) + self.assertTrue(np.array_equal(np.array(out), expected)) + def test_cummax_cummin_nan(self): nan = float("nan") cases = [ From b16ef730f2b53fe2e355b82adb7a33add158a57e Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 19 Aug 2026 10:00:50 +0900 Subject: [PATCH 2/2] nit --- mlx/backend/cpu/scan.cpp | 5 +---- python/tests/test_ops.py | 2 +- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/mlx/backend/cpu/scan.cpp b/mlx/backend/cpu/scan.cpp index 23151c8e02..93e67825e2 100644 --- a/mlx/backend/cpu/scan.cpp +++ b/mlx/backend/cpu/scan.cpp @@ -163,10 +163,7 @@ void scan_op( const Op& op, U init) { if (in.flags().row_contiguous) { - // A size one axis can carry any stride and still be row contiguous, so - // the strided count below would come out zero. The scan is elementwise - // in that case, which is what the contiguous path with a stride of one - // computes. + // 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(), diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index ccbc9390dc..ff88ce7414 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2642,7 +2642,7 @@ def test_scan_size_one_axis(self): 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, stream=mx.cpu) + out = getattr(mx, op)(a, axis=0) self.assertTrue(np.array_equal(np.array(out), expected)) def test_cummax_cummin_nan(self):