fix: keep single-element strided slice region round-tripping

mlx/ops.cpp

@@ -726,12 +726,13 @@ array flip(const array& a, StreamOrDevice s /* = {} */) {
 namespace {
 
 inline auto
-normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) {
+normalize_slice(const Shape& shape, Shape& start, Shape& stop, Shape& strides) {
   // - Start indices are normalized
   // - End indices are unchanged as -1 means something different
   //   pre-normalization (the end of the axis) versus post normalization (the
   //   position left of 0).
-  // - Any strides corresponding to singleton dimension are set to 1
+  // - Any strides corresponding to singleton dimension are set to 1 (and
+  //   stop is narrowed so the stored triple still round-trips)
 
   Shape out_shape(shape.size());
   bool has_neg_strides = false;
@@ -772,9 +773,25 @@ normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) {
 
       out_shape[i] = (ed - start[i] + strides[i] - 1) / strides[i];
     }
-    // Simplify the stride if it's unused
+    // Simplify the stride if it's unused. Keep the stored triple
+    // round-tripping so consumers that re-derive the region from it
+    // (Slice's vjp and vmap) select the same single element the forward
+    // pass read, instead of the whole half-open span or an empty region.
     if (out_shape[i] == 1) {
-      strides[i] = 1;
+      if (strides[i] < 0) {
+        if (start[i] > 0) {
+          strides[i] = -1;
+          stop[i] = start[i] - 1;
+        } else {
+          // start == 0 cannot be expressed as a single-element
+          // negative-stride span, use the equivalent unit stride.
+          strides[i] = 1;
+          stop[i] = 1;
+        }
+      } else {
+        strides[i] = 1;
+        stop[i] = start[i] + 1;
+      }
     }
   }
 

python/tests/test_autograd.py

@@ -1261,6 +1261,48 @@ def fun(a, b):
         expected[4:-5:-2] = tan_b
         self.assertTrue(mx.allclose(grad, expected))
 
+    def test_slice_grads_single_element(self):
+        # A strided slice selecting exactly one element must place each
+        # cotangent at the position the forward pass read it from, not over
+        # the whole half-open span (regression: the stride was collapsed to 1
+        # without narrowing stop, so the vjp re-derived a wider region).
+        def grad_of(fun, x):
+            return mx.grad(fun)(x)
+
+        # positive stride, first element of a length-2 axis
+        x = mx.zeros((2, 3))
+        Wa = mx.ones((1, 3))
+        Wb = mx.ones((1, 3)) * 2
+        fun = lambda z: (z[0::2] * Wa).sum() + (z[1::2] * Wb).sum()
+        g = grad_of(fun, x)
+        self.assertTrue(mx.allclose(g, mx.array([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]])))
+
+        # positive stride, single element in a longer span
+        x = mx.zeros((8, 2))
+        W = mx.ones((1, 2))
+        g = grad_of(lambda z: (z[2:5:3] * W).sum(), x)
+        expected = mx.zeros((8, 2))
+        expected[2] = W
+        self.assertTrue(mx.allclose(g, expected))
+
+        # negative stride, single element
+        x = mx.zeros((2, 2))
+        W = mx.ones((1, 2))
+        g = grad_of(lambda z: (z[::-2] * W).sum(), x)
+        expected = mx.zeros((2, 2))
+        expected[1] = W
+        self.assertTrue(mx.allclose(g, expected))
+
+        # negative stride with start == 0 (cannot be expressed as a
+        # single-element negative-stride span; the unit-stride fallback must
+        # still land the cotangent at index 0)
+        x = mx.zeros((3, 2))
+        W = mx.ones((1, 2))
+        g = grad_of(lambda z: (z[0::-1] * W).sum(), x)
+        expected = mx.zeros((3, 2))
+        expected[0] = W
+        self.assertTrue(mx.allclose(g, expected))
+
     def test_leaks(self):
         for transform in [
             mx.grad,

python/tests/test_vmap.py

@@ -219,6 +219,19 @@ def test_vmap_indexing(self):
         )
         self.assertTrue(mx.array_equal(out, expected))
 
+    def test_vmap_strided_slice_single_element(self):
+        # A strided slice selecting exactly one element must vmap to the
+        # batched version of what the un-batched slice returns (regression:
+        # the stored stride was collapsed to 1 without narrowing stop, so
+        # re-deriving the region returned every element in the span).
+        x = mx.arange(24, dtype=mx.float32).reshape(4, 2, 3)
+        out = mx.vmap(lambda t: t[0::2])(x)
+        self.assertTrue(mx.array_equal(out, x[:, 0:1]))
+
+        y = mx.arange(48, dtype=mx.float32).reshape(4, 4, 3)
+        out = mx.vmap(lambda t: t[0::2])(y)
+        self.assertTrue(mx.array_equal(out, y[:, 0:4:2]))
+
     def test_vmap_reduce(self):
         a = mx.ones((5, 5), mx.int32)
         out = mx.vmap(lambda x: x.sum())(a)