diff --git a/python/src/indexing.cpp b/python/src/indexing.cpp index 59b0655d..564c4cb4 100644 --- a/python/src/indexing.cpp +++ b/python/src/indexing.cpp @@ -15,12 +15,20 @@ bool is_none_slice(const nb::slice& in_slice) { nb::getattr(in_slice, "step").is_none()); } +int safe_to_int32(nb::object obj) { + auto val = nb::cast(nb::cast(obj)); + if (val > INT32_MAX || val < INT32_MIN) { + throw std::invalid_argument("Slice indices must be 32-bit integers."); + } + return static_cast(val); +} + int get_slice_int(nb::object obj, int default_val) { if (!obj.is_none()) { if (!nb::isinstance(obj)) { throw std::invalid_argument("Slice indices must be integers or None."); } - return nb::cast(nb::cast(obj)); + return safe_to_int32(obj); } return default_val; } @@ -45,7 +53,7 @@ void get_slice_params( } mx::array get_int_index(nb::object idx, int axis_size) { - int idx_ = nb::cast(idx); + int idx_ = safe_to_int32(idx); idx_ = (idx_ < 0) ? idx_ + axis_size : idx_; return mx::array(idx_, mx::uint32); diff --git a/python/tests/test_array.py b/python/tests/test_array.py index 89ca5ffd..4efed9da 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -2127,6 +2127,13 @@ class TestArray(mlx_tests.MLXTestCase): self.assertEqual(x.imag.item(), 1.0) self.assertEqual(x.real.item(), 1.0) + def test_large_indices(self): + x = mx.array([0, 1, 2]) + with self.assertRaises(ValueError): + x[: 2**32] + with self.assertRaises(ValueError): + x[2**32] + if __name__ == "__main__": mlx_tests.MLXTestRunner()