mlx icon indicating copy to clipboard operation
mlx copied to clipboard

Fix deep recursion with siblings

Open awni opened this issue 1 year ago • 1 comments

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)

awni avatar Oct 06 '24 15:10 awni

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.

awni avatar Oct 07 '24 03:10 awni