Fix deep recursion with siblings (#1462)

* fix recursion with siblings

* fix

* add test

* increase tol

python/tests/test_array.py

@@ -1868,6 +1868,33 @@ class TestArray(mlx_tests.MLXTestCase):
         with self.assertRaises(ValueError):
             int(a)
 
+    def test_deep_graphs(self):
+        # The following tests should simply run cleanly without a segfault or
+        # crash due to exceeding recursion depth limits.
+
+        # 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(100_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)
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_conv.py

@@ -902,7 +902,7 @@ class TestConv(mlx_tests.MLXTestCase):
             dw10 = (cotan[1::s, :-1:s] * x).sum()
             dw11 = (cotan[1::s, 1::s] * x).sum()
             expected = mx.array([[dw00, dw01], [dw10, dw11]])
-            self.assertTrue(mx.allclose(dw, expected))
+            self.assertTrue(mx.allclose(dw, expected, rtol=1e-5, atol=1e-5))
 
     def test_conv_groups_grad(self):
         def fn(x, w):