zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-04 15:04:40 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix array is_available race cases (#1468)

Browse Source
This commit is contained in:
Awni Hannun
2024-10-07 19:13:50 -07:00
committed by GitHub
parent 9b12093739
commit 3274c6a087
9 changed files with 77 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/tests/test_eval.py
View File

@@ -51,11 +51,13 @@ class TestEval(mlx_tests.MLXTestCase):
self.assertTrue(mx.array_equal(z, mx.array([4, 8, 12])))
def test_async_eval_twice(self):
x = mx.array(1) + mx.array(1) + mx.array(1)
mx.async_eval(x)
y = x + 1
mx.async_eval(y)
self.assertEqual(x.item(), 3)
for _ in range(1000):
x = mx.array(1) + mx.array(1) + mx.array(1)
mx.async_eval(x)
y = x + 1
mx.async_eval(y)
self.assertEqual(x.item(), 3)
self.assertEqual(y.item(), 4)
def test_async_eval_in_trace(self):
def fun(x):