zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-31 16:21:27 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Revert "More buffer donation in some cases (#1858)" (#1863)

Browse Source 
This reverts commit d274ae77f2.
This commit is contained in:
Awni Hannun
2025-02-13 14:21:44 -08:00
committed by GitHub
parent 5cd97f7ffe
commit 428f589364
2 changed files with 4 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
python/tests/test_eval.py
View File

@@ -174,25 +174,6 @@ class TestEval(mlx_tests.MLXTestCase):
post = mx.metal.get_peak_memory()
self.assertEqual(pre, post)
def test_donation_multiple_inputs(self):
def fun(its, x, y):
for _ in range(its):
a = x + y # y should donate
b = x + a # x should donate
x, y = a, b
return x, y
x = mx.zeros((128, 128))
y = mx.zeros((128, 128))
mx.metal.reset_peak_memory()
a, b = fun(2, x, y)
mx.eval(a, b)
mem2 = mx.metal.get_peak_memory()
a, b = fun(10, x, y)
mx.eval(a, b)
mem10 = mx.metal.get_peak_memory()
self.assertEqual(mem2, mem10)
if __name__ == "__main__":
unittest.main()