From 855195828e043737e5552152fa6451f19392d23d Mon Sep 17 00:00:00 2001 From: Ishaan Date: Tue, 15 Sep 2026 00:03:58 -0400 Subject: [PATCH] Clamp slice bounds when a slice is combined with an array index When a slice is combined with an array index, the slice bounds are resolved for negative values but never clamped into [0, axis_size], so out-of-range bounds are passed to arange() verbatim: start = (start < 0) ? start + src.shape(i) : start; A single adjustment is not enough: for a 4-element axis, -100 becomes -96, and arange(-96, 100) yields 196 indices into that axis. Consequences on a (4, 5) array with idx = [0, 1]: a[0:100, idx] -> (100, 2) of repeated data, NumPy gives (4, 2) a[-100:4, idx] -> (100, 2), NumPy gives (4, 2) a[-100:-50, idx] -> (50, 2), NumPy gives (0, 2) a[-100:4, idx] = 0 -> writes 1 row, NumPy writes 4 a[100:-100:-1, idx] -> SIGBUS, NumPy gives (4, 2) The last case is an out-of-bounds read reachable from ordinary Python. Slicing past the end is idiomatic and safe elsewhere in Python, so the first case in particular is easy to hit by accident, and none of these raise. Add adjust_slice_bounds(), which mirrors CPython's PySlice_AdjustIndices (clamping to [0, n] for positive strides and [-1, n-1] for negative ones), and use it in both mlx_gather_nd and mlx_scatter_args_nd, which carried the same partial adjustment. Plain slicing and slice-plus-integer indexing were already correct and are untouched. Tested with a differential sweep against NumPy over 11 x 11 x 7 start/stop/step combinations for both __getitem__ and __setitem__. --- python/src/indexing.cpp | 35 +++++++++++++++++++++++++++++------ python/tests/test_array.py | 27 +++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/python/src/indexing.cpp b/python/src/indexing.cpp index 5bf4677505..5238948b13 100644 --- a/python/src/indexing.cpp +++ b/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(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(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); diff --git a/python/tests/test_array.py b/python/tests/test_array.py index e03803631a..0c2e8bd9ba 100644 --- a/python/tests/test_array.py +++ b/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 = [