mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-19 00:04:41 +08:00
shapeless compile in docs and partially shapeless reshape (#1742)
This commit is contained in:
@@ -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()
|
||||
|
Reference in New Issue
Block a user