zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-02 17:28:12 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

move memory APIs into top level mlx.core (#1982)

Browse Source
This commit is contained in:
Awni Hannun
2025-03-21 07:25:12 -07:00
committed by GitHub
parent 65a38c452b
commit 4e1994e9d7
25 changed files with 418 additions and 323 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/tests/test_eval.py
View File

@@ -118,9 +118,9 @@ class TestEval(mlx_tests.MLXTestCase):
out = mx.vjp(fn, (x,), (y,))
out = mx.vjp(fn, (x,), (y,))
if mx.metal.is_available():
peak_mem = mx.metal.get_peak_memory()
peak_mem = mx.get_peak_memory()
out = mx.vjp(fn, (x,), (y,))
self.assertEqual(peak_mem, mx.metal.get_peak_memory())
self.assertEqual(peak_mem, mx.get_peak_memory())
def test_async_eval_with_multiple_streams(self):
x = mx.array([1.0])
@@ -151,11 +151,11 @@ class TestEval(mlx_tests.MLXTestCase):
x = mx.zeros((4096, 4096))
mx.eval(x)
pre = mx.metal.get_peak_memory()
pre = mx.get_peak_memory()
out = fun(x)
del x
mx.eval(out)
post = mx.metal.get_peak_memory()
post = mx.get_peak_memory()
self.assertEqual(pre, post)
def fun(x):
@@ -167,11 +167,11 @@ class TestEval(mlx_tests.MLXTestCase):
x = mx.zeros((4096 * 4096,))
mx.eval(x)
pre = mx.metal.get_peak_memory()
pre = mx.get_peak_memory()
out = fun(x)
del x
mx.eval(out)
post = mx.metal.get_peak_memory()
post = mx.get_peak_memory()
self.assertEqual(pre, post)
@unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
@@ -187,7 +187,7 @@ class TestEval(mlx_tests.MLXTestCase):
s1 = mx.default_stream(mx.gpu)
s2 = mx.new_stream(mx.gpu)
old_limit = mx.metal.set_memory_limit(1000)
old_limit = mx.set_memory_limit(1000)
x = mx.ones((512, 512), stream=s2)
for _ in range(80):
@@ -195,7 +195,7 @@ class TestEval(mlx_tests.MLXTestCase):
y = mx.abs(x, stream=s2)
z = mx.abs(y, stream=s2)
mx.eval(z)
mx.metal.set_memory_limit(old_limit)
mx.set_memory_limit(old_limit)
if __name__ == "__main__":