From 4ba705df5b01c35267bc240e6350b2b1c5076a5c Mon Sep 17 00:00:00 2001 From: devangpratap <115096812+devangpratap@users.noreply.github.com> Date: Tue, 8 Sep 2026 23:32:01 -0400 Subject: [PATCH] Raise IndexError for out of bounds axes An out of bounds axis is an indexing error, but most axis checks threw std::invalid_argument, which nanobind surfaces as ValueError. The reduction path already threw std::out_of_range (IndexError), so mx.sum(x, axis=5) and mx.expand_dims(x, 5) disagreed on the same mistake. numpy raises AxisError for both, which subclasses ValueError and IndexError. Switch the bounds checks to std::out_of_range. Most of them go through normalize_axis_index, the rest are ad hoc checks in ops, fft, random and the vmap bindings. Errors that are not about bounds, duplicate axes, axis count mismatches and shape mismatches, stay ValueError. Closes #4428 --- mlx/fft.cpp | 4 +-- mlx/ops.cpp | 20 +++++------ mlx/random.cpp | 2 +- mlx/utils.cpp | 2 +- python/src/transforms.cpp | 4 +-- python/tests/test_fft.py | 4 +-- python/tests/test_ops.py | 71 ++++++++++++++++++++++++++++++++------- python/tests/test_vmap.py | 8 ++--- tests/fft_tests.cpp | 11 +++--- tests/ops_tests.cpp | 8 ++--- 10 files changed, 90 insertions(+), 44 deletions(-) diff --git a/mlx/fft.cpp b/mlx/fft.cpp index 06860a0e3a..55596472ad 100644 --- a/mlx/fft.cpp +++ b/mlx/fft.cpp @@ -61,7 +61,7 @@ array fft_impl( std::ostringstream msg; msg << "[fftn] Invalid axis received for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // In the following shape manipulations there are three cases to consider: @@ -262,7 +262,7 @@ array fftshift_impl( std::ostringstream msg; msg << "[" << name << "] Invalid axis " << ax << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Match NumPy's implementation int shift = a.shape(axis) / 2; diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 1a41688d28..6ee4690a36 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -482,7 +482,7 @@ array unflatten( std::ostringstream msg; msg << "[unflatten] Invalid axes " << ax << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } size_t size = 1; @@ -817,7 +817,7 @@ void normalize_dynamic_slice_inputs( std::ostringstream msg; msg << prefix << " Invalid axis " << ax << " for array with dimension " << a.ndim() << "."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } ax = new_ax; } @@ -1770,7 +1770,7 @@ array transpose( std::ostringstream msg; msg << "[transpose] Invalid axis (" << ax << ") for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } if (shape[ax] != 0) { throw std::invalid_argument("[transpose] Repeat axes not allowed."); @@ -2333,7 +2333,7 @@ array mean( std::ostringstream msg; msg << "[mean] axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } } auto dtype = at_least_float(a.dtype()); @@ -2367,7 +2367,7 @@ array median( std::ostringstream msg; msg << "[median] axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } set_axes.insert(axis < 0 ? axis + ndim : axis); } @@ -3677,7 +3677,7 @@ array take( std::ostringstream msg; msg << "[take] Received invalid axis " << axis << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Check for valid take @@ -3722,7 +3722,7 @@ array take(const array& a, int index, int axis, StreamOrDevice s /* = {} */) { std::ostringstream msg; msg << "[take] Received invalid axis " << axis << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Check for valid take @@ -4249,7 +4249,7 @@ array diff( int ndim = static_cast(a.ndim()); int ax = axis < 0 ? axis + ndim : axis; if (ax < 0 || ax >= ndim) { - throw std::invalid_argument("[diff] Axis is out of bounds for the array."); + throw std::out_of_range("[diff] Axis is out of bounds for the array."); } if (n < 0) { throw std::invalid_argument("[diff] Order `n` must be non-negative."); @@ -5796,7 +5796,7 @@ array vecdot( } int ax = axis < 0 ? axis + a.ndim() : axis; if (ax < 0 || ax >= a.ndim()) { - throw std::invalid_argument("[vecdot] axis is out of bounds."); + throw std::out_of_range("[vecdot] axis is out of bounds."); } if (axis < 0 ? axis + b.ndim() != ax : axis >= b.ndim()) { throw std::invalid_argument("[vecdot] axis is out of bounds."); @@ -6677,7 +6677,7 @@ array roll( std::ostringstream msg; msg << "[roll] Invalid axis " << axes[i] << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } auto sh = shift[i]; diff --git a/mlx/random.cpp b/mlx/random.cpp index 2932b3eaf8..9a49853ff6 100644 --- a/mlx/random.cpp +++ b/mlx/random.cpp @@ -383,7 +383,7 @@ int get_valid_axis(int axis, int ndim) { std::ostringstream msg; msg << "[categorical] Invalid axis " << axis << " for logits with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } return ax; } diff --git a/mlx/utils.cpp b/mlx/utils.cpp index 4e29e8847b..07130fa690 100644 --- a/mlx/utils.cpp +++ b/mlx/utils.cpp @@ -179,7 +179,7 @@ int normalize_axis_index( std::ostringstream msg; msg << msg_prefix << "Axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } return axis < 0 ? axis + ndim : axis; } diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp index e33cde9bce..6c2299b7f8 100644 --- a/python/src/transforms.cpp +++ b/python/src/transforms.cpp @@ -334,7 +334,7 @@ auto py_vmap( msg << "[vmap] Invalid" << (output_axes ? " output " : " ") << "vectorization axis " << axis << " for array with shape " << x.shape(); - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } flat_axes.push_back(axis); } else if (nb::isinstance(inputs[1])) { @@ -351,7 +351,7 @@ auto py_vmap( msg << "[vmap] Invalid" << (output_axes ? " output " : " ") << "vectorization axis " << axis << " for array with shape " << x.shape(); - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } flat_axes.push_back(axis); } else if (l.size() == 1 && l[0].is_none()) { diff --git a/python/tests/test_fft.py b/python/tests/test_fft.py index 1f96aad566..60f7610fe6 100644 --- a/python/tests/test_fft.py +++ b/python/tests/test_fft.py @@ -385,9 +385,9 @@ def test_ifftshift(self): def test_fftshift_errors(self): # Test invalid axes x = mx.array(np.random.rand(4, 4).astype(np.float32)) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.fft.fftshift(x, axes=[2]) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.fft.fftshift(x, axes=[-3]) # Test empty array diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index d5568610e8..5aeacb3a2a 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -977,7 +977,7 @@ def test_median(self): with self.assertRaises(ValueError): mx.median(x, axis=0) x = mx.array([0, 1, 2, 3, 4]) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.median(x, axis=(0, 1)) with self.assertRaises(ValueError): mx.median(x, axis=(0, 0)) @@ -1562,9 +1562,9 @@ def test_along_axis_invalid_axis(self): values = mx.ones(a.shape, dtype=a.dtype) for ax in [3, 4, 100, -4, -5, -100]: - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.take_along_axis(a, idx, axis=ax) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.put_along_axis(a, idx, values, axis=ax) # Valid negative axes still work @@ -1576,7 +1576,7 @@ def test_cross_invalid_axis(self): a = mx.array([1.0, 2.0, 3.0]) b = mx.array([4.0, 5.0, 6.0]) for ax in [1, 2, -2, -50]: - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.linalg.cross(a, b, axis=ax) def test_put_along_axis(self): @@ -1646,7 +1646,7 @@ def test_split(self): self.assertEqual(y.tolist(), [[3, 4]]) self.assertEqual(z.tolist(), [[5, 6]]) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.split(a, 3, axis=2) a = mx.arange(8) @@ -1668,7 +1668,7 @@ def test_flip(self): b_np = np.array([1, 2, 3, 4]) self.assertTrue(np.array_equal(mx.flip(mx.array(b_np)), np.flip(b_np))) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.flip(a, axis=2) def test_unstack(self): @@ -1693,7 +1693,7 @@ def test_unstack(self): # stack is the inverse of unstack. self.assertTrue(mx.array_equal(mx.stack(mx.unstack(a, axis=1), axis=1), a)) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.unstack(a, axis=2) def test_split_invalid_num_splits(self): @@ -2579,7 +2579,7 @@ def test_scan_invalid_axis(self): for op in ["cumsum", "cumprod", "cummax", "cummin", "logcumsumexp"]: mxop = getattr(mx, op) for ax in [3, 4, 100, -4, -5, -100]: - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mxop(a, axis=ax) # Valid negative axes still work and agree with the positive one @@ -2783,7 +2783,7 @@ def test_diff(self): self.assertEqual(mx.diff(m, axis=0).tolist(), [[-1, 2, 0]]) self.assertEqual(mx.diff(m, axis=1).tolist(), [[2, 3], [5, 1]]) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.diff(a, axis=1) def test_squeeze_expand(self): @@ -2807,19 +2807,64 @@ def test_squeeze_expand_invalid_axes(self): # Out of bounds negative axes must raise instead of wrapping around a = mx.zeros(()) self.assertEqual(mx.expand_dims(a, (-2, -1)).shape, (1, 1)) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.expand_dims(a, (-3, -2)) a = mx.zeros((2, 2)) for axes in [(-5, -4), (-6, 0), (0, 5)]: - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.expand_dims(a, axes) a = mx.zeros((1, 1, 1)) for axes in [(-4,), (-5, 0), (0, 4)]: - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.squeeze(a, axes) + def test_out_of_bounds_axis_raises_index_error(self): + # An out of bounds axis is an indexing error, not a value error, so + # every op agrees on IndexError (numpy raises AxisError, which is a + # subclass of both). + a = mx.zeros((2, 3)) + cases = [ + lambda: mx.sum(a, axis=5), + lambda: mx.mean(a, axis=5), + lambda: mx.median(a, axis=5), + lambda: mx.expand_dims(a, 5), + lambda: mx.squeeze(mx.zeros((1, 1)), 5), + lambda: mx.moveaxis(a, 5, 0), + lambda: mx.swapaxes(a, 5, 0), + lambda: mx.transpose(a, (5, 0)), + lambda: mx.roll(a, 1, 5), + lambda: mx.diagonal(a, 0, 5, 1), + lambda: mx.take(a, mx.array([0]), axis=5), + lambda: mx.vecdot(a, a, axis=5), + lambda: mx.diff(a, axis=5), + lambda: mx.take_along_axis(a, mx.zeros((2, 3), mx.int32), axis=5), + lambda: mx.cumsum(a, axis=5), + lambda: mx.argsort(a, axis=5), + lambda: mx.split(a, 2, axis=5), + lambda: mx.flip(a, axis=5), + lambda: mx.fft.fft(a, axis=5), + lambda: mx.fft.fftshift(a, axes=[5]), + lambda: mx.random.categorical(a, axis=5), + lambda: mx.unflatten(a, 5, (1, -1)), + lambda: mx.slice(a, mx.array([0]), [5], (1,)), + lambda: mx.slice_update(a, mx.zeros((1, 1)), mx.array([0]), [5]), + lambda: mx.vmap(mx.exp, in_axes=5)(a), + ] + for case in cases: + with self.assertRaises(IndexError): + case() + + # Duplicate or mismatched axes stay ValueError, they are not a + # bounds problem. + with self.assertRaises(ValueError): + mx.transpose(a, (0, 0)) + with self.assertRaises(ValueError): + mx.transpose(a, (0,)) + with self.assertRaises(ValueError): + mx.sum(a, axis=(0, 0)) + def test_sort(self): shape = (6, 4, 10) tests = product( @@ -3388,7 +3433,7 @@ def test_vecdot(self): with self.assertRaises(ValueError): mx.vecdot(mx.array(1), mx.array([1])) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.vecdot(mx.array([1, 2]), mx.array([1, 2]), axis=1) with self.assertRaises(ValueError): mx.vecdot(mx.array([1, 2]), mx.array([1])) diff --git a/python/tests/test_vmap.py b/python/tests/test_vmap.py index b050a80384..7538b559db 100644 --- a/python/tests/test_vmap.py +++ b/python/tests/test_vmap.py @@ -9,8 +9,8 @@ class TestVmap(mlx_tests.MLXTestCase): def test_basics(self): - # Can't vmap over scalars - with self.assertRaises(ValueError): + # Can't vmap over scalars, axis 0 is out of bounds for a 0d array + with self.assertRaises(IndexError): mx.vmap(mx.exp)(mx.array(1.0)) # Invalid input @@ -21,13 +21,13 @@ def test_basics(self): with self.assertRaises(ValueError): mx.vmap(mx.exp, in_axes="hello")(mx.array([0, 1])) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.vmap(mx.exp, in_axes=2)(mx.array([0, 1])) with self.assertRaises(ValueError): mx.vmap(mx.exp, out_axes="hello")(mx.array([0, 1])) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.vmap(mx.exp, out_axes=2)(mx.array([0, 1])) def test_unary(self): diff --git a/tests/fft_tests.cpp b/tests/fft_tests.cpp index 7dc1b6e666..6adc4c1d4c 100644 --- a/tests/fft_tests.cpp +++ b/tests/fft_tests.cpp @@ -134,6 +134,9 @@ TEST_CASE("test real ffts") { TEST_CASE("test fftn") { auto x = zeros({5, 5, 5}); CHECK_THROWS_AS(fft::fftn(x, {}, {0, 3}), std::invalid_argument); + // Matching n/axes sizes so the bounds check is what trips. + CHECK_THROWS_AS(fft::fftn(x, {5, 5}, {0, 3}), std::out_of_range); + CHECK_THROWS_AS(fft::fftn(x, {5, 5}, {0, -4}), std::out_of_range); CHECK_THROWS_AS(fft::fftn(x, {}, {0, -4}), std::invalid_argument); CHECK_THROWS_AS(fft::fftn(x, {}, {0, 0}), std::invalid_argument); CHECK_THROWS_AS(fft::fftn(x, {5, 5, 5}, {0}), std::invalid_argument); @@ -388,8 +391,8 @@ TEST_CASE("test fftshift and ifftshift") { CHECK(array_equal(y, expected).item()); // Test error cases - CHECK_THROWS_AS(fft::fftshift(x, {3}), std::invalid_argument); - CHECK_THROWS_AS(fft::fftshift(x, {-5}), std::invalid_argument); - CHECK_THROWS_AS(fft::ifftshift(x, {3}), std::invalid_argument); - CHECK_THROWS_AS(fft::ifftshift(x, {-5}), std::invalid_argument); + CHECK_THROWS_AS(fft::fftshift(x, {3}), std::out_of_range); + CHECK_THROWS_AS(fft::fftshift(x, {-5}), std::out_of_range); + CHECK_THROWS_AS(fft::ifftshift(x, {3}), std::out_of_range); + CHECK_THROWS_AS(fft::ifftshift(x, {-5}), std::out_of_range); } diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 7c36bca762..9c95e9b5e9 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -706,7 +706,7 @@ TEST_CASE("test transpose") { CHECK_EQ(y.shape(), Shape{1}); CHECK_EQ(y.item(), 1); - CHECK_THROWS_AS(transpose(x, {1}), std::invalid_argument); + CHECK_THROWS_AS(transpose(x, {1}), std::out_of_range); CHECK_THROWS_AS(transpose(x, {0, 0}), std::invalid_argument); // Works with empty array @@ -4627,11 +4627,9 @@ TEST_CASE("test pad with an axes subset") { // An axis outside the array is rejected rather than indexed. for (auto mode : all) { CHECK_THROWS_AS( - pad(x, {5}, Shape{1}, Shape{1}, array(0.0f), mode), - std::invalid_argument); + pad(x, {5}, Shape{1}, Shape{1}, array(0.0f), mode), std::out_of_range); CHECK_THROWS_AS( - pad(x, {-5}, Shape{1}, Shape{1}, array(0.0f), mode), - std::invalid_argument); + pad(x, {-5}, Shape{1}, Shape{1}, array(0.0f), mode), std::out_of_range); } }