Repository navigation
Fix compile_fuse broadcast split aliasing bug - #3166
Merged
angeloskath merged 2 commits intoFeb 27, 2026
Merged
angeloskath merged 2 commits into
angeloskath merged 2 commits into
Conversation
split_one creates a copy of a broadcast array and moves cache parents to point to the copy by modifying their input vectors. However, the array reference `a` in the recurse lambda is aliased through a parent's input vector. When split_one modifies that parent's input, `a` silently becomes the copy, losing the reference to the original. When the split copy then fails to be fused (e.g. hitting max_compile_depth or max_compile_arrays), no code path adds the original to input_set, so compile_replace crashes with unordered_map::at on the second (cached) call. Fix: save the original array before calling split_one. If the split copy fails to be fused, use merge_one to undo the split and treat the original as an input. This preserves the broadcast fusion optimization from commit 2c11d10 while fixing the crash. Fixes ml-explore#3165
angeloskath
approved these changes
Feb 27, 2026
angeloskath
left a comment
angeloskath
left a comment
Member
There was a problem hiding this comment.
Thanks for finding the bug. I simplified the test and solution and will merge after the tests pass.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #3165 —
mx.compilecrashes withunordered_map::at: key not foundon computation graphs with >~100 operations.split_onemodifies a parent's input vector, silently replacing theconst array& areference in therecurselambda (which is aliased through that vector). When the split copy fails to fuse, the original broadcast is orphaned — never added toinput_set.array orig(a)beforesplit_one. If the copy isn't fused, undo withmerge_oneand use the original as an input. Preserves the broadcast fusion optimization from Split broadcast so it is always fused in compile #2318.test_compile_large_graph_with_broadcasts— builds a ~150-op graph (Gaussian scoring with 15 observations), verifies both cached execution and correctness vs uncompiled.Test plan
test_compile_large_graph_with_broadcastspassestest_shared_broadcaststill passes (optimization preserved)🤖 Generated with Claude Code