Clearer error when shape dimension overflows int32 (#3425)

Co-authored-by: Kanishk <kanishk.chores@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

python/tests/test_array.py

@@ -767,10 +767,13 @@ class TestArray(mlx_tests.MLXTestCase):
 
     def test_array_np_shape_dim_check(self):
         a_npy = np.empty(2**31, dtype=np.bool_)
-        with self.assertRaises(ValueError) as e:
+        with self.assertRaises(OverflowError) as e:
             mx.array(a_npy)
         self.assertEqual(
-            str(e.exception), "Shape dimension falls outside supported `int` range."
+            str(e.exception),
+            "Shape dimension 2147483648 is outside the supported range "
+            "[-2147483648, 2147483647]. MLX currently uses 32-bit integers "
+            "for shape dimensions.",
         )
 
     def test_dtype_promotion(self):

python/tests/test_ops.py

@@ -92,6 +92,52 @@ class TestOps(mlx_tests.MLXTestCase):
             self.assertEqual(y.dtype, t)
             self.assertTrue(mx.array_equal(y, x))
 
+    def test_shape_overflow_error(self):
+        # Shape dimensions that don't fit in int32 should raise a clear
+        # OverflowError that names the offending value, rather than a generic
+        # "incompatible function arguments" TypeError. The overflow check
+        # lives in the mx::Shape type caster, so it applies to every op that
+        # takes a shape. See issue #2681.
+        too_big = 2**31
+
+        # Array creation ops — also exercise the scalar shape path.
+        for ctor in (mx.zeros, mx.ones):
+            with self.assertRaises(OverflowError) as cm:
+                ctor(too_big)
+            self.assertIn(str(too_big), str(cm.exception))
+            with self.assertRaises(OverflowError) as cm:
+                ctor([too_big])
+            self.assertIn(str(too_big), str(cm.exception))
+
+        with self.assertRaises(OverflowError) as cm:
+            mx.full(too_big, 0.0)
+        self.assertIn(str(too_big), str(cm.exception))
+        with self.assertRaises(OverflowError) as cm:
+            mx.full([too_big], 0.0)
+        self.assertIn(str(too_big), str(cm.exception))
+
+        # Other shape-taking ops should surface the same clean error.
+        a = mx.zeros(4)
+        with self.assertRaises(OverflowError) as cm:
+            mx.reshape(a, [too_big])
+        self.assertIn(str(too_big), str(cm.exception))
+
+        with self.assertRaises(OverflowError) as cm:
+            mx.broadcast_to(a, [too_big, 1])
+        self.assertIn(str(too_big), str(cm.exception))
+
+        # Negative overflow (< int32 min) is caught too.
+        too_negative = -(2**31) - 1
+        with self.assertRaises(OverflowError) as cm:
+            mx.zeros([too_negative])
+        self.assertIn(str(too_negative), str(cm.exception))
+
+        # Shapes that fit in int32 still go through unchanged.
+        self.assertEqual(mx.zeros(4).shape, (4,))
+        self.assertEqual(mx.zeros((2, 3)).shape, (2, 3))
+        self.assertEqual(mx.ones([2, 3]).shape, (2, 3))
+        self.assertEqual(mx.full((2, 3), 1.5).tolist(), [[1.5] * 3] * 2)
+
     def test_scalar_inputs(self):
         # Check combinations of python types
         a = mx.add(False, True)