Raise IndexError for out of bounds axes

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;

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<int>(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];

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

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

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<nb::tuple>(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()) {

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

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]))

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):

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<bool>());
 
   // 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);
 }

tests/ops_tests.cpp

@@ -706,7 +706,7 @@ TEST_CASE("test transpose") {
   CHECK_EQ(y.shape(), Shape{1});
   CHECK_EQ(y.item<int>(), 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);
   }
 }