ml-explore / mlx

MLX: An array framework for Apple silicon
https://ml-explore.github.io/mlx/
MIT License
14.83k stars 845 forks source link

Fix leak for multi-output primitives which are never detached #1059

Closed awni closed 2 weeks ago

awni commented 2 weeks ago

Closes: https://github.com/ml-explore/mlx-examples/issues/738

No issues destroying very deep graphs, the following still works:

a = mx.array(1.0)
for _ in range(1000000):
    a = mx.abs(a)
awni commented 2 weeks ago

@angeloskath I added some comments in the destructor to help since I don't have any great ideas to simplify it. Lmk what you think. Should be working now.