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 = [