From c8536f5248bf06e20a8ed1706c57a6687d2e7c64 Mon Sep 17 00:00:00 2001 From: Robert Johansson Date: Fri, 27 Feb 2026 03:17:28 +0100 Subject: [PATCH] Fix compile_fuse broadcast split aliasing bug (#3166) Co-authored-by: Angelos Katharopoulos --- mlx/compile.cpp | 3 ++- python/tests/test_compile.py | 22 ++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/mlx/compile.cpp b/mlx/compile.cpp index ca5f0699..4af68d0c 100644 --- a/mlx/compile.cpp +++ b/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 { diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index d64c057f..7db471cc 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -1049,6 +1049,28 @@ class TestCompile(mlx_tests.MLXTestCase): 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():