Bound the strided scan to the size of its output

mlx/backend/cuda/scan.cu

@@ -458,7 +458,10 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& encoder = cu::get_command_encoder(s);
 
-  if (in.flags().contiguous && in.strides()[axis_] != 0) {
+  // The strided kernel writes out[i * stride + j] for i < shape[axis] and
+  // j < stride, so the scanned axis has to fit inside the allocation.
+  if (in.flags().contiguous && in.strides()[axis_] != 0 &&
+      in.shape(axis_) * in.strides()[axis_] <= in.data_size()) {
     if (in.is_donatable() && in.itemsize() == out.itemsize()) {
       out.copy_shared_buffer(in);
     } else {

mlx/backend/metal/scan.cpp

@@ -117,7 +117,10 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
 
   auto in = inputs[0];
-  if (in.flags().contiguous && in.strides()[axis_] != 0) {
+  // The strided kernel writes out[i * stride + j] for i < shape[axis] and
+  // j < stride, so the scanned axis has to fit inside the allocation.
+  if (in.flags().contiguous && in.strides()[axis_] != 0 &&
+      in.shape(axis_) * in.strides()[axis_] <= in.data_size()) {
     if (in.is_donatable() && in.itemsize() == out.itemsize()) {
       out.copy_shared_buffer(in);
     } else {

python/tests/test_ops.py

@@ -2500,6 +2500,20 @@ def fn(its):
         mem4 = mx.get_peak_memory()
         self.assertEqual(mem2, mem4)
 
+    def test_scan_does_not_write_past_the_output(self):
+        # A size one axis can carry any stride, and the strided kernel bounds
+        # its writes by that stride rather than by the size of the output.
+        for n, off in ((64, 63), (512, 400)):
+            with self.subTest(n=n, off=off):
+                base = mx.arange(1, n + 1, dtype=mx.float32).reshape(1, n)
+                canaries = [mx.zeros((16,)) + i for i in range(12)]
+                mx.eval(base, *canaries)
+                before = [np.array(x) for x in [base] + canaries]
+                mx.eval(mx.cumsum(base[:, off:], axis=0))
+                after = [np.array(x) for x in [base] + canaries]
+                for b, a in zip(before, after):
+                    self.assertTrue(np.array_equal(b, a))
+
     def test_cummax_cummin_nan(self):
         nan = float("nan")
         cases = [

tests/gpu_tests.cpp

@@ -712,3 +712,15 @@ TEST_CASE("test layer norm vjp bias grad race") {
   }
   CHECK(worst <= 1e-5);
 }
+
+TEST_CASE("test scan output holds the strided kernel writes") {
+  // A size one axis can carry a padded stride. The strided scan addresses
+  // out[i * stride + j] for i < shape[axis] and j < stride, so the output has
+  // to hold shape[axis] * stride elements. This slice overruns by one element,
+  // the tightest case.
+  auto a = slice(reshape(arange(64.0, float32), {1, 64}), {0, 1}, {1, 64});
+  eval(a);
+  auto out = cumsum(a, 0);
+  eval(out);
+  CHECK_LE(out.shape(0) * out.strides()[0], out.data_size());
+}