L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Clearer error when shape dimension overflows int32 (#3425)

Browse Source 
Co-authored-by: Kanishk <[email protected]>
Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
serenposh
2026-05-05 09:53:36 +09:00
committed by GitHub
co-authored by Kanishk Claude Opus 4.7
parent b43965925f
commit 1fdd4e23c2
6 changed files with 122 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_ops.py
+46
View File
@@ -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)