Fix: move subgraph initializers to main graph when inlining If node branches

When constant folding inlines an If node's branch (because the condition is
constant), initializers from the subgraph were not moved to the main graph.
This caused them to be lost (e.g., the 'axes' initializer for a Squeeze node
inside an If branch).

Fix: add _move_initializers_to_graph() helper that moves initializers from
the subgraph to the main graph with name uniqueness handling, and call it
from if_op() before returning the Replacement.

Adds test: test_fold_if_cond_with_subgraph_initializer

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

Author: Copilot

=0.1.16

@@ -561,6 +561,21 @@ def size(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     return op.Constant(value_int=size)
 
 
+def _move_initializers_to_graph(src: ir.Graph, dst: ir.Graph) -> None:
+    """Move all initializers from src graph to dst graph, ensuring name uniqueness."""
+    counter: dict[str, int] = {}
+    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)
+
+
 @register("If")
 def if_op(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     cond_input = _get_input(node, 0)
@@ -598,7 +613,11 @@ def rename(name):
             # Avoid name collision.
             sub_node.name = f"{node.name}_{sub_node.name}"
 
-        # TODO: we should handle initializers as well!
+        # Move initializers from the subgraph to the main graph to avoid losing them.
+        main_graph = node.graph
+        if main_graph is not None:
+            _move_initializers_to_graph(graph, main_graph)
+
         return Replacement(formal_outs, graph_nodes)
     return None
 

onnxscript/optimizer/_constant_folding.py

@@ -6,6 +6,8 @@
 
 import numpy as np
 import onnx
+import onnx.helper
+import onnx.numpy_helper
 import parameterized
 
 import onnxscript.optimizer as optimizer
@@ -130,6 +132,72 @@ def test_fold_if_cond(self):
         self.assertEqual(optimized.graph[0].outputs[0].name, "z")
         self.assertEqual(optimized.graph[0].op_type, "Mul")
 
+    def _make_if_model_with_subgraph_initializer(self) -> ir.Model:
+        """Build a model where the then_branch of an If node has an initializer."""
+        # Build then_branch: Squeeze(x, axes) where axes=[0] is an initializer
+        axes_init = onnx.numpy_helper.from_array(np.array([0], dtype=np.int64), name="axes")
+        squeeze_node = onnx.helper.make_node(
+            "Squeeze", inputs=["x", "axes"], outputs=["z"], name="Squeeze_0"
+        )
+        then_branch = onnx.helper.make_graph(
+            [squeeze_node],
+            "then_branch",
+            [],
+            [onnx.helper.make_tensor_value_info("z", onnx.TensorProto.FLOAT, [16, 16])],
+            initializer=[axes_init],
+        )
+        # Build else_branch: Identity(x)
+        identity_node = onnx.helper.make_node(
+            "Identity", inputs=["x"], outputs=["z"], name="Identity_0"
+        )
+        else_branch = onnx.helper.make_graph(
+            [identity_node],
+            "else_branch",
+            [],
+            [onnx.helper.make_tensor_value_info("z", onnx.TensorProto.FLOAT, [1, 16, 16])],
+        )
+        # Build main graph with a constant True condition
+        one_node = onnx.helper.make_node(
+            "Constant",
+            inputs=[],
+            outputs=["one"],
+            value=onnx.helper.make_tensor("one", onnx.TensorProto.BOOL, [], [True]),
+        )
+        if_node = onnx.helper.make_node(
+            "If",
+            inputs=["one"],
+            outputs=["z"],
+            then_branch=then_branch,
+            else_branch=else_branch,
+            name="If_1",
+        )
+        main_graph = onnx.helper.make_graph(
+            [one_node, if_node],
+            "main",
+            [onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 16, 16])],
+            [onnx.helper.make_tensor_value_info("z", onnx.TensorProto.FLOAT, [16, 16])],
+        )
+        proto = onnx.helper.make_model(
+            main_graph, opset_imports=[onnx.helper.make_opsetid("", 17)]
+        )
+        proto.ir_version = 8
+        return ir.serde.deserialize_model(proto)
+
+    def test_fold_if_cond_with_subgraph_initializer(self):
+        """Test that initializers in inlined If branches are moved to the main graph."""
+        model = self._make_if_model_with_subgraph_initializer()
+        self.assertIn("axes", model.graph[1].attributes["then_branch"].as_graph().initializers)
+        self.assertNotIn("axes", model.graph.initializers)
+
+        optimized = self._fold(model)
+
+        # The If node should be inlined; the axes initializer must be in the main graph
+        self.assertIn("axes", optimized.graph.initializers)
+        np.testing.assert_array_equal(
+            optimized.graph.initializers["axes"].const_value.numpy(),
+            np.array([0], dtype=np.int64),
+        )
+
     def test_fold_inside_if_branch(self):
         model = """
             <ir_version: 7, opset_import: [ "" : 17]>