mlx
mlx copied to clipboard
Fix deep recursion with siblings
The non-recursive destructor didn't catch two cases:
- Destruction of arrays that have duplicate inputs
- Destruction of arrays with siblings
So this fixes that.
Close https://github.com/ml-explore/mlx-swift/issues/139
Added some tests to the long running CI:
def deep_graph():
# Deep graph destroyed without eval
x = mx.array([1.0, 2.0])
for _ in range(100_000):
x = mx.sin(x)
del x
# Duplicate input deep graph destroyed without eval
x = mx.array([1.0, 2.0])
for _ in range(500_000):
x = x + x
# Deep graph with siblings destroyed without eval
x = mx.array([1, 2])
for _ in range(100_000):
x = mx.concatenate(mx.split(x, 2))
del x
# Deep graph with eval
x = mx.array([1.0, 2.0])
for _ in range(100_000):
x = mx.sin(x)
mx.eval(x)
The tests seem fast enough for the normal CI as well don't they?
I think you're right. I'll add them here instead.