Clamp slice bounds when a slice is combined with an array index

python/src/indexing.cpp

@@ -75,6 +75,33 @@ void get_slice_params(
       nb::getattr(in_slice, "stop"), strides < 0 ? -axis_size - 1 : axis_size);
 }
 
+// Resolve negative bounds and clamp into range, mirroring CPython's
+// PySlice_AdjustIndices. Needed wherever a slice is expanded into an explicit
+// arange, since out-of-range bounds would otherwise be used verbatim.
+void adjust_slice_bounds(
+    mx::ShapeElem& start,
+    mx::ShapeElem& end,
+    mx::ShapeElem stride,
+    int axis_size) {
+  if (start < 0) {
+    start += axis_size;
+    if (start < 0) {
+      start = (stride < 0) ? -1 : 0;
+    }
+  } else if (start >= axis_size) {
+    start = (stride < 0) ? axis_size - 1 : axis_size;
+  }
+
+  if (end < 0) {
+    end += axis_size;
+    if (end < 0) {
+      end = (stride < 0) ? -1 : 0;
+    }
+  } else if (end >= axis_size) {
+    end = (stride < 0) ? axis_size - 1 : axis_size;
+  }
+}
+
 mx::array get_int_index(nb::object idx, int axis_size) {
   int idx_ = safe_to_int32(idx);
   idx_ = (idx_ < 0) ? idx_ + axis_size : idx_;
@@ -155,9 +182,7 @@ mx::array mlx_gather_nd(
       get_slice_params(
           start, end, stride, nb::cast<nb::slice>(idx), src.shape(i));
 
-      // Handle negative indices
-      start = (start < 0) ? start + src.shape(i) : start;
-      end = (end < 0) ? end + src.shape(i) : end;
+      adjust_slice_bounds(start, end, stride, src.shape(i));
 
       gather_indices.push_back(arange(start, end, stride, mx::uint32));
       num_slices++;
@@ -686,9 +711,7 @@ mlx_scatter_args_nd(
       get_slice_params(
           start, end, stride, nb::cast<nb::slice>(pyidx), axis_size);
 
-      // Handle negative indices
-      start = (start < 0) ? start + axis_size : start;
-      end = (end < 0) ? end + axis_size : end;
+      adjust_slice_bounds(start, end, stride, axis_size);
 
       mx::Shape idx_shape(idx_ndim, 1);
 

python/tests/test_array.py

@@ -1882,6 +1882,33 @@ def test_slice_negative_step(self):
         b_mx = a_mx[::-1, ::-3, ::-2]
         self.assertTrue(np.array_equal(b_np, b_mx))
 
+    def test_slice_bounds_with_array_index(self):
+        # Out-of-range slice bounds must be clamped the same way NumPy clamps
+        # them when the slice is combined with an array index.
+        a_np = np.arange(20, dtype=np.int32).reshape(4, 5)
+        a_mx = mx.array(a_np)
+        idx_np = np.array([0, 1])
+        idx_mx = mx.array([0, 1], dtype=mx.uint32)
+
+        bounds = [None, -100, -6, -4, -1, 0, 1, 3, 4, 6, 100]
+        steps = [None, 1, 2, 3, -1, -2, -3]
+
+        for start in bounds:
+            for stop in bounds:
+                for step in steps:
+                    s = slice(start, stop, step)
+                    with self.subTest(start=start, stop=stop, step=step):
+                        self.assertTrue(
+                            np.array_equal(a_np[s, idx_np], a_mx[s, idx_mx])
+                        )
+
+                        # Same clamping applies when assigning through the slice
+                        u_np = a_np.copy()
+                        u_mx = mx.array(a_np)
+                        u_np[s, idx_np] = 0
+                        u_mx[s, idx_mx] = 0
+                        self.assertTrue(np.array_equal(u_np, u_mx))
+
     def test_api(self):
         x = mx.array(np.random.rand(10, 10, 10))
         ops = [