L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add clear_streams API for cleanup before exit (#3395)

Browse Source
This commit is contained in:
Cheng
2026-04-14 18:41:32 +09:00
committed by GitHub
parent 520cea2bec
commit 435f0b6cdb
10 changed files with 53 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_threads.py
+3 -1
View File
@@ -7,7 +7,7 @@ import mlx.core as mx
import mlx_tests
class TestReduce(mlx_tests.MLXTestCase):
class TestThreads(mlx_tests.MLXTestCase):
def test_threadlocal_stream(self):
test_stream = mx.new_stream(mx.default_device())
@@ -16,6 +16,7 @@ class TestReduce(mlx_tests.MLXTestCase):
with mx.stream(test_stream):
x = mx.arange(10)
mx.eval(2 * x)
mx.clear_streams()
t1 = threading.Thread(target=test_failure)
t2 = threading.Thread(target=test_failure)
@@ -31,6 +32,7 @@ class TestReduce(mlx_tests.MLXTestCase):
x = mx.arange(10)
mx.eval(2 * x)
self.assertEqual(x.tolist(), list(range(10)))
mx.clear_streams()
t1 = threading.Thread(target=test_success)
t2 = threading.Thread(target=test_success)