fix(optimizer): move subgraph initializers to main graph when inlining If branches#2887
fix(optimizer): move subgraph initializers to main graph when inlining If branches#2887jaytiwarihub wants to merge 1 commit intomicrosoft:mainfrom
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Fixes a correctness issue in constant folding where inlined If branches could leave behind subgraph initializers, producing invalid ONNX graphs with dangling value references.
Changes:
- Add
_move_initializers_to_graph(src, dst)and call it fromif_opbefore returning the inliningReplacement - Add regression tests covering subgraph-initializer migration and a “name collision” scenario
- Minor formatting cleanups in tests and comments/log messages
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 9 comments.
| File | Description |
|---|---|
| onnxscript/optimizer/_constant_folding.py | Migrates initializers from an inlined If branch subgraph into the main graph and does some minor doc/log cleanup. |
| onnxscript/optimizer/_constant_folding_test.py | Adds regression tests for the initializer-migration bug and adjusts formatting of several existing test definitions/comments. |
Comments suppressed due to low confidence (1)
onnxscript/optimizer/_constant_folding.py:1457
- There’s trailing whitespace on the return line. Also, removing
# type: ignore[return-value]may reintroduce static type-check failures ifInPlacePass.__call__isn’t typed to returnFoldConstantsResult. If the repo runs mypy/pyright, consider either restoring the targeted ignore or adjusting the pass/call typing sofolder_pass(model)is correctly typed without ignores.
def fold_constants(
| return op.Constant(value_int=size) | ||
|
|
||
|
|
||
| def _move_initializers_to_graph(src: ir.Graph, dst: ir.Graph) -> None: |
There was a problem hiding this comment.
This helper guarantees uniqueness only against dst.initializers, but ONNX value names must be unique across the whole graph namespace (e.g., node outputs). If name collides with an existing non-initializer value name in dst, register_initializer may raise or the graph may become invalid. Also, popping from src.initializers before successfully registering can irreversibly drop the initializer if an exception occurs. Consider: (1) building a used_names set for the destination graph covering all existing value names (not just initializers), (2) choosing new_name against that set, and (3) only removing from src.initializers after the destination registration succeeds.
| for name in list(src.initializers): | ||
| initializer = src.initializers.pop(name) | ||
| # Ensure name uniqueness in the destination graph | ||
| new_name = name | ||
| while new_name in dst.initializers: | ||
| counter[name] = counter.get(name, 0) + 1 | ||
| new_name = f"{name}_{counter[name]}" | ||
| if new_name != name: | ||
| initializer.name = new_name | ||
| dst.register_initializer(initializer) |
There was a problem hiding this comment.
This helper guarantees uniqueness only against dst.initializers, but ONNX value names must be unique across the whole graph namespace (e.g., node outputs). If name collides with an existing non-initializer value name in dst, register_initializer may raise or the graph may become invalid. Also, popping from src.initializers before successfully registering can irreversibly drop the initializer if an exception occurs. Consider: (1) building a used_names set for the destination graph covering all existing value names (not just initializers), (2) choosing new_name against that set, and (3) only removing from src.initializers after the destination registration succeeds.
| if main_graph is not None: | ||
| _move_initializers_to_graph(graph, main_graph) | ||
|
|
||
| # TODO: we should handle initializers as well! |
There was a problem hiding this comment.
The TODO on line 623 is now stale/contradictory because the function is explicitly handling initializers. Please remove or update that TODO to avoid misleading future readers.
| # TODO: we should handle initializers as well! |
| should_fold: An optional function that takes a node and returns True if | ||
| the node should be considered for folding. | ||
| The function should return True/False value to indicate if this particular | ||
| The function should return True/False value to indicate if this particular |
There was a problem hiding this comment.
In the class docstring, the sentence on line 994 is mis-indented relative to the should_fold: attribute description, which makes the rendered doc confusing (it reads like a new attribute). Re-indent it so it’s clearly part of should_fold’s description.
| The function should return True/False value to indicate if this particular | |
| The function should return True/False value to indicate if this particular |
| should_fold=should_fold, | ||
| ) | ||
| return folder_pass(model) # type: ignore[return-value] | ||
| return folder_pass(model) No newline at end of file |
There was a problem hiding this comment.
There’s trailing whitespace on the return line. Also, removing # type: ignore[return-value] may reintroduce static type-check failures if InPlacePass.__call__ isn’t typed to return FoldConstantsResult. If the repo runs mypy/pyright, consider either restoring the targeted ignore or adjusting the pass/call typing so folder_pass(model) is correctly typed without ignores.
| return folder_pass(model) | |
| return typing.cast(FoldConstantsResult, folder_pass(model)) |
| def test_fold_if_cond_with_subgraph_initializer_name_collision(self): | ||
| """Subgraph initializer names that clash with main-graph names get a unique suffix.""" |
There was a problem hiding this comment.
This test currently asserts only that it 'must not crash' and that the final model passes onnx.checker.check_model, but it does not verify that a name collision actually occurred or that suffixing happened (and the chosen graph text may not force a collision depending on what gets folded into initializers). To make this a robust regression test for the collision logic, assert that: (1) the main graph already contains an initializer with the colliding name before inlining, (2) after inlining there are two distinct initializers (original + suffixed), and (3) the inlined nodes reference the suffixed initializer (not the original).
| def test_static_split_to_sequence_with_scalar_split_and_squence_at_is_folded_as_split( | ||
| self, | ||
| ): | ||
| def test_static_split_to_sequence_with_scalar_split_and_squence_at_is_folded_as_split(self): |
There was a problem hiding this comment.
The word 'squence' appears to be a misspelling of 'sequence' in these updated test names. Since these lines were modified in this PR, it’s a good opportunity to correct the spelling for clarity and searchability.
| def test_static_split_to_sequence_with_scalar_split_and_squence_at_is_folded_as_split(self): | |
| def test_static_split_to_sequence_with_scalar_split_and_sequence_at_is_folded_as_split(self): |
| def test_static_split_to_sequence_with_list_split_and_squence_at_is_folded_as_split( | ||
| self, | ||
| ): | ||
| def test_static_split_to_sequence_with_list_split_and_squence_at_is_folded_as_split(self): |
There was a problem hiding this comment.
The word 'squence' appears to be a misspelling of 'sequence' in these updated test names. Since these lines were modified in this PR, it’s a good opportunity to correct the spelling for clarity and searchability.
| def test_static_split_to_sequence_with_list_split_no_keepdims_and_squence_at_is_folded_as_split_with_squeeze( | ||
| self, | ||
| ): | ||
| def test_static_split_to_sequence_with_list_split_no_keepdims_and_squence_at_is_folded_as_split_with_squeeze(self): |
There was a problem hiding this comment.
The word 'squence' appears to be a misspelling of 'sequence' in these updated test names. Since these lines were modified in this PR, it’s a good opportunity to correct the spelling for clarity and searchability.
|
@microsoft-github-policy-service agree |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| should_fold=should_fold, | ||
| ) | ||
| return folder_pass(model) # type: ignore[return-value] | ||
| return folder_pass(model) No newline at end of file |
| should_fold=should_fold, | ||
| ) | ||
| return folder_pass(model) # type: ignore[return-value] | ||
| return folder_pass(model) No newline at end of file |
| should_fold=should_fold, | ||
| ) | ||
| return folder_pass(model) # type: ignore[return-value] | ||
| return folder_pass(model) No newline at end of file |
|
|
||
| if __name__ == "__main__": | ||
| unittest.main() | ||
| unittest.main() No newline at end of file |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2887 +/- ##
==========================================
+ Coverage 72.02% 72.26% +0.23%
==========================================
Files 239 241 +2
Lines 29309 29690 +381
Branches 2882 2918 +36
==========================================
+ Hits 21110 21455 +345
- Misses 7219 7238 +19
- Partials 980 997 +17 ☔ View full report in Codecov by Sentry. |
Fixes #2810
When
FoldConstantsPassinlines anIfbranch (because its condition is a known constant), initializers belonging to the branch subgraph were silently abandoned. This caused downstream nodes that referenced those initializers — most visiblySqueezenodes whoseaxestensor had been folded in a prior pass — to have dangling inputs, producing an invalid ONNX model that fails schema validation.Cause
The
if_oppartial evaluator correctly moves the branch's nodes into the main graph viaReplacement, but it never touched the branch's initializers. Initializers are stored on their.Graphobject itself, not on individual nodes, so they were invisible to the node-migration logic. After inlining, those values existed nowhere in the main graph.Fix
Added a helper
_move_initializers_to_graph(src, dst)that migrates all initializers from a subgraph to the destination graph before theReplacementis returned, with name-collision handling.Testing
Added two new tests to
_constant_folding_test.py:test_fold_if_cond_with_subgraph_initializer: Reproduces the exact two-pass scenario from the bug report. The first pass folds a constant inside the branch into a subgraph initializer; the second pass inlines the branch and verifies the initializer is successfully migrated to the main graph and that the model passesonnx.checker.check_model.test_fold_if_cond_with_subgraph_initializer_name_collision: Verifies that name collisions between subgraph and main-graph initializers are handled gracefully via unique suffixing, preventing data loss or crashes during the migration.