| 14 | |
| 15 | class TestCompile(mlx_tests.MLXTestCase): |
| 16 | def test_simple_compile(self): |
| 17 | def fun(x, y): |
| 18 | return x + y |
| 19 | |
| 20 | compiled_fn = mx.compile(fun) |
| 21 | compiled_fn = mx.compile(fun) |
| 22 | x = mx.array(1.0) |
| 23 | y = mx.array(1.0) |
| 24 | out = compiled_fn(x, y) |
| 25 | self.assertEqual(out.item(), 2.0) |
| 26 | |
| 27 | # Try again |
| 28 | out = compiled_fn(x, y) |
| 29 | self.assertEqual(out.item(), 2.0) |
| 30 | |
| 31 | # Change sizes |
| 32 | x = mx.array([1.0, 2.0]) |
| 33 | out = compiled_fn(x, y) |
| 34 | self.assertTrue(mx.array_equal(out, mx.array([2.0, 3.0]))) |
| 35 | |
| 36 | y = mx.array([1.0, 2.0]) |
| 37 | out = compiled_fn(x, y) |
| 38 | self.assertTrue(mx.array_equal(out, mx.array([2.0, 4.0]))) |
| 39 | |
| 40 | # Change types |
| 41 | x = mx.array([1, 2], mx.int32) |
| 42 | y = mx.array([1, 2], mx.int32) |
| 43 | out = compiled_fn(x, y) |
| 44 | self.assertEqual(out.dtype, mx.int32) |
| 45 | self.assertTrue(mx.array_equal(out, mx.array([2, 4]))) |
| 46 | |
| 47 | def test_compile_grad(self): |
| 48 | def loss_fn(x): |