diff --git a/python/src/indexing.cpp b/python/src/indexing.cpp index 1dce537880..2e421db518 100644 --- a/python/src/indexing.cpp +++ b/python/src/indexing.cpp @@ -151,9 +151,23 @@ mx::array mlx_gather_nd( get_slice_params( start, end, stride, nb::cast(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++; diff --git a/python/tests/test_array.py b/python/tests/test_array.py index 80459fabe5..a178d4fa64 100644 --- a/python/tests/test_array.py +++ b/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)