Clamp slice bounds in mixed fancy/slice indexing

python/src/indexing.cpp

@@ -151,9 +151,23 @@ mx::array mlx_gather_nd(
       get_slice_params(
           start, end, stride, nb::cast<nb::slice>(idx), src.shape(i));
 
+      auto axis_size = src.shape(i);
       // Handle negative indices
-      start = (start < 0) ? start + src.shape(i) : start;
-      end = (end < 0) ? end + src.shape(i) : end;
+      start = (start < 0) ? start + axis_size : start;
+      end = (end < 0) ? end + axis_size : end;
+
+      // Clamp to the valid range for this axis, matching the behavior of
+      // mlx::core::slice / normalize_slice, so out-of-range or heavily
+      // negative bounds don't produce an incorrectly sized/valued gather.
+      if (stride < 0) {
+        start = std::min(start, axis_size - 1);
+        end = std::max(end, mx::ShapeElem{-1});
+        end = std::min(end, start);
+      } else {
+        start = std::max(mx::ShapeElem{0}, std::min(start, axis_size));
+        end = std::max(mx::ShapeElem{0}, std::min(end, axis_size));
+        end = std::max(end, start);
+      }
 
       gather_indices.push_back(arange(start, end, stride, mx::uint32));
       num_slices++;

python/tests/test_array.py

@@ -1267,6 +1267,37 @@ def check_slices(arr_np, *idx_np):
         a_mlx = mx.array(a_np)
         self.assertTrue(np.array_equal(a_np[2:-1, 0], np.array(a_mlx[2:-1, 0])))
 
+    def test_indexing_mixed_slice_array_out_of_bounds(self):
+        # Regression test: when a slice is combined with an array/int index
+        # (e.g. a[start:stop, idx_array]), the slice bounds were adjusted for
+        # negative indices but never clamped into the valid [0, axis_size]
+        # range before being passed to arange(). Out-of-range or heavily
+        # negative slice bounds therefore produced arrays of the wrong shape
+        # containing bogus repeated/garbage data instead of matching NumPy's
+        # clamping behavior.
+        a_npy = np.arange(20, dtype=np.int32).reshape(4, 5)
+        a_mlx = mx.array(a_npy)
+        idx_npy = np.array([0, 1], dtype=np.uint32)
+        idx_mlx = mx.array(idx_npy)
+
+        # Large negative start, in-range stop
+        out_mlx = a_mlx[-100:4, idx_mlx]
+        out_npy = a_npy[-100:4, idx_npy]
+        self.assertEqual(out_mlx.shape, out_npy.shape)
+        self.assertTrue(np.array_equal(np.asarray(out_mlx), out_npy))
+
+        # Out-of-range stop
+        out_mlx = a_mlx[0:200, idx_mlx]
+        out_npy = a_npy[0:200, idx_npy]
+        self.assertEqual(out_mlx.shape, out_npy.shape)
+        self.assertTrue(np.array_equal(np.asarray(out_mlx), out_npy))
+
+        # Very large negative start (would previously allocate a huge array)
+        out_mlx = a_mlx[-(10**9) : 4, idx_mlx]
+        out_npy = a_npy[-(10**9) : 4, idx_npy]
+        self.assertEqual(out_mlx.shape, out_npy.shape)
+        self.assertTrue(np.array_equal(np.asarray(out_mlx), out_npy))
+
     def test_indexing_grad(self):
         x = mx.array([[1, 2], [3, 4]]).astype(mx.float32)
         ind = mx.array([0, 1, 0]).astype(mx.float32)