| 33 | self.assertEqual(x.item(), 3) |
| 34 | |
| 35 | def test_async_eval(self): |
| 36 | x = mx.array(1) + mx.array(1) + mx.array(1) |
| 37 | mx.async_eval(x) |
| 38 | self.assertEqual(x.item(), 3) |
| 39 | |
| 40 | # It should be safe to call eval on the array which has been async |
| 41 | # eval'ed |
| 42 | x = mx.array(1) + mx.array(1) + mx.array(1) |
| 43 | self.assertEqual(x.item(), 3) |
| 44 | |
| 45 | x = mx.array([1, 2, 3]) |
| 46 | y = 2 * x |
| 47 | mx.async_eval(y) |
| 48 | z = 2 * y |
| 49 | mx.async_eval(z) |
| 50 | self.assertTrue(mx.array_equal(y, mx.array([2, 4, 6]))) |
| 51 | self.assertTrue(mx.array_equal(z, mx.array([4, 8, 12]))) |
| 52 | |
| 53 | def test_async_eval_twice(self): |
| 54 | for _ in range(1000): |