recompile when shapeless is different (#1776)

2025-09-04 15:04:40 +08:00 · 2025-01-20 21:07:10 -08:00
parent a8666a757a
commit 90532b1f37
2 changed files with 11 additions and 2 deletions
--- a/python/tests/test_compile.py
+++ b/python/tests/test_compile.py
@@ -675,6 +675,10 @@ class TestCompile(mlx_tests.MLXTestCase):
        def mean(x):
            return mx.mean(x, keepdims=True)

+        cfun = mx.compile(mean)
+        out = cfun(mx.ones((5, 5)))
+        self.assertTrue(mx.allclose(out, mx.array(1.0)))
+
        cmean = mx.compile(mean, shapeless=True)

        x = mx.ones(2)