shapeless compile in docs and partially shapeless reshape (#1742)

This commit is contained in:
Awni Hannun
2025-01-02 16:24:42 -08:00
committed by GitHub
parent a64a8dfe45
commit ae69cb15e9
6 changed files with 137 additions and 36 deletions

View File

@@ -830,6 +830,25 @@ class TestCompile(mlx_tests.MLXTestCase):
a = mx.array([0.0, 1.0, 2.0, 3.0, 4.0])
self.assertTrue(mx.allclose(cfun(a), fun(a)))
def test_shapeless_compile_with_reshape(self):
def fun(x):
return x.reshape(x.shape[0] * x.shape[1], -1)
compiled_fun = mx.compile(fun, shapeless=True)
x = mx.zeros(shape=(2, 3, 4))
out = compiled_fun(x)
self.assertEqual(out.shape, (6, 4))
x = mx.zeros(shape=(2, 3, 8))
out = compiled_fun(x)
self.assertEqual(out.shape, (6, 8))
x = mx.zeros(shape=(5, 5, 5))
with self.assertRaises(ValueError):
compiled_fun(x)
if __name__ == "__main__":
unittest.main()