Fix compile_fuse broadcast split aliasing bug (ml-explore#3166)

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

Author: robert-johansson

mlx/compile.cpp

@@ -856,7 +856,8 @@ void compile_fuse(
       // are not fusable except for broadcast which we can split to avoid
       // stopping fusion
       if (!all_parents_in) {
-        if (a.has_primitive() && is_broadcast(a.primitive())) {
+        if (a.has_primitive() && is_broadcast(a.primitive()) &&
+            input_set.size() < max_compile_arrays) {
           array b = split_one(a, parents_map, cache);
           recurse(b, depth, s, shape);
         } else {

python/tests/test_compile.py

@@ -1049,6 +1049,28 @@ def fun(x, y, z):
         self.assertTrue(mx.allclose(d[0], d_hat[0]))
         self.assertTrue(mx.allclose(d[1], d_hat[1]))
 
+    def test_compile_large_graph_with_broadcasts(self):
+        N = 20
+        _as = [mx.array(2 * i, dtype=mx.float32) for i in range(N)]
+        _bs = [mx.array(i, dtype=mx.float32) for i in range(N)]
+        _c = mx.array(0.0)
+        x = mx.random.normal((2, 2))
+
+        def f(x):
+            y = 0
+            for i in range(N):
+                y = y + _as[i] * x * _bs[i] * _c
+            return y
+
+        ref = f(x)
+        mx.eval(ref)
+        f = mx.compile(f)
+        for i in range(2):
+            y = f(x)
+            mx.eval(y)
+
+        self.assertTrue(mx.allclose(y, ref))
+
     def test_wrap_compiled(self):
         @mx.compile
         def inner():