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); } }