From c282d4c0ca56e2d73351644ae400a17fb5742508 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 25 Feb 2026 02:14:10 +0000 Subject: [PATCH 1/4] Initial plan From c3afbcee38d5d2fd2d8da8ebec40115d448e0287 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 25 Feb 2026 02:22:51 +0000 Subject: [PATCH 2/4] Initial plan for fixing constant folding with If node initializers Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> --- =0.1.16 | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 =0.1.16 diff --git a/=0.1.16 b/=0.1.16 new file mode 100644 index 0000000000..3f5aa80d37 --- /dev/null +++ b/=0.1.16 @@ -0,0 +1,32 @@ +Defaulting to user installation because normal site-packages is not writeable +Collecting onnx_ir + Downloading onnx_ir-0.2.0-py3-none-any.whl.metadata (3.3 kB) +Collecting numpy (from onnx_ir) + Downloading numpy-2.4.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (6.6 kB) +Collecting onnx>=1.16 (from onnx_ir) + Downloading onnx-1.20.1-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (8.4 kB) +Requirement already satisfied: typing_extensions>=4.10 in /usr/lib/python3/dist-packages (from onnx_ir) (4.10.0) +Collecting ml_dtypes>=0.5.0 (from onnx_ir) + Downloading ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (8.9 kB) +Collecting sympy>=1.13 (from onnx_ir) + Downloading sympy-1.14.0-py3-none-any.whl.metadata (12 kB) +Collecting protobuf>=4.25.1 (from onnx>=1.16->onnx_ir) + Downloading protobuf-6.33.5-cp39-abi3-manylinux2014_x86_64.whl.metadata (593 bytes) +Collecting mpmath<1.4,>=1.1.0 (from sympy>=1.13->onnx_ir) + Downloading mpmath-1.3.0-py3-none-any.whl.metadata (8.6 kB) +Downloading onnx_ir-0.2.0-py3-none-any.whl (164 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 164.1/164.1 kB 49.2 MB/s eta 0:00:00 +Downloading ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (5.0 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 5.0/5.0 MB 158.1 MB/s eta 0:00:00 +Downloading numpy-2.4.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (16.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 16.6/16.6 MB 173.5 MB/s eta 0:00:00 +Downloading onnx-1.20.1-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (17.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 17.5/17.5 MB 101.3 MB/s eta 0:00:00 +Downloading sympy-1.14.0-py3-none-any.whl (6.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.3/6.3 MB 172.9 MB/s eta 0:00:00 +Downloading mpmath-1.3.0-py3-none-any.whl (536 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 536.2/536.2 kB 94.6 MB/s eta 0:00:00 +Downloading protobuf-6.33.5-cp39-abi3-manylinux2014_x86_64.whl (323 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 323.5/323.5 kB 79.6 MB/s eta 0:00:00 +Installing collected packages: mpmath, sympy, protobuf, numpy, ml_dtypes, onnx, onnx_ir +Successfully installed ml_dtypes-0.5.4 mpmath-1.3.0 numpy-2.4.2 onnx-1.20.1 onnx_ir-0.2.0 protobuf-6.33.5 sympy-1.14.0 From c5803399307cc3a89027df4781976adb02aa095e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 25 Feb 2026 02:27:39 +0000 Subject: [PATCH 3/4] 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> --- =0.1.16 | 32 --------- onnxscript/optimizer/_constant_folding.py | 21 +++++- .../optimizer/_constant_folding_test.py | 68 +++++++++++++++++++ 3 files changed, 88 insertions(+), 33 deletions(-) delete mode 100644 =0.1.16 diff --git a/=0.1.16 b/=0.1.16 deleted file mode 100644 index 3f5aa80d37..0000000000 --- a/=0.1.16 +++ /dev/null @@ -1,32 +0,0 @@ -Defaulting to user installation because normal site-packages is not writeable -Collecting onnx_ir - Downloading onnx_ir-0.2.0-py3-none-any.whl.metadata (3.3 kB) -Collecting numpy (from onnx_ir) - Downloading numpy-2.4.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (6.6 kB) -Collecting onnx>=1.16 (from onnx_ir) - Downloading onnx-1.20.1-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (8.4 kB) -Requirement already satisfied: typing_extensions>=4.10 in /usr/lib/python3/dist-packages (from onnx_ir) (4.10.0) -Collecting ml_dtypes>=0.5.0 (from onnx_ir) - Downloading ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (8.9 kB) -Collecting sympy>=1.13 (from onnx_ir) - Downloading sympy-1.14.0-py3-none-any.whl.metadata (12 kB) -Collecting protobuf>=4.25.1 (from onnx>=1.16->onnx_ir) - Downloading protobuf-6.33.5-cp39-abi3-manylinux2014_x86_64.whl.metadata (593 bytes) -Collecting mpmath<1.4,>=1.1.0 (from sympy>=1.13->onnx_ir) - Downloading mpmath-1.3.0-py3-none-any.whl.metadata (8.6 kB) -Downloading onnx_ir-0.2.0-py3-none-any.whl (164 kB) - ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 164.1/164.1 kB 49.2 MB/s eta 0:00:00 -Downloading ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (5.0 MB) - ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 5.0/5.0 MB 158.1 MB/s eta 0:00:00 -Downloading numpy-2.4.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (16.6 MB) - ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 16.6/16.6 MB 173.5 MB/s eta 0:00:00 -Downloading onnx-1.20.1-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (17.5 MB) - ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 17.5/17.5 MB 101.3 MB/s eta 0:00:00 -Downloading sympy-1.14.0-py3-none-any.whl (6.3 MB) - ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.3/6.3 MB 172.9 MB/s eta 0:00:00 -Downloading mpmath-1.3.0-py3-none-any.whl (536 kB) - ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 536.2/536.2 kB 94.6 MB/s eta 0:00:00 -Downloading protobuf-6.33.5-cp39-abi3-manylinux2014_x86_64.whl (323 kB) - ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 323.5/323.5 kB 79.6 MB/s eta 0:00:00 -Installing collected packages: mpmath, sympy, protobuf, numpy, ml_dtypes, onnx, onnx_ir -Successfully installed ml_dtypes-0.5.4 mpmath-1.3.0 numpy-2.4.2 onnx-1.20.1 onnx_ir-0.2.0 protobuf-6.33.5 sympy-1.14.0 diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 574ddd8aef..26375cc4dc 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -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 diff --git a/onnxscript/optimizer/_constant_folding_test.py b/onnxscript/optimizer/_constant_folding_test.py index 080af9c2f3..740321b289 100644 --- a/onnxscript/optimizer/_constant_folding_test.py +++ b/onnxscript/optimizer/_constant_folding_test.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 = """ From 58755be1487594bac19013fece5d4f0b898c0408 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 25 Feb 2026 02:40:00 +0000 Subject: [PATCH 4/4] Refactor: rewrite test to use ONNX text format as requested Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> --- .../optimizer/_constant_folding_test.py | 103 +++++++----------- 1 file changed, 37 insertions(+), 66 deletions(-) diff --git a/onnxscript/optimizer/_constant_folding_test.py b/onnxscript/optimizer/_constant_folding_test.py index 740321b289..75a454a39d 100644 --- a/onnxscript/optimizer/_constant_folding_test.py +++ b/onnxscript/optimizer/_constant_folding_test.py @@ -6,8 +6,6 @@ import numpy as np import onnx -import onnx.helper -import onnx.numpy_helper import parameterized import onnxscript.optimizer as optimizer @@ -132,71 +130,44 @@ 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), - ) + """If branch initializers should be moved to the main graph when the branch is inlined.""" + # A model with a non-constant condition; constants inside the then_branch will + # be folded into subgraph initializers on the first fold pass. + model = ir.from_onnx_text(""" + + agraph (float[16, 16] x, bool cond) => (float[16, 16] z) { + two = Constant () + three = Constant () + z = If (cond) < + then_branch = then_graph () => (then_z) { + temp = Add (two, three) + then_z = Mul (temp, x) + }, + else_branch = else_graph () => (else_z) { + else_z = Identity (x) + } + > + } + """) + # First fold: 'temp = Add(2.0, 3.0)' gets folded into a subgraph initializer. + _constant_folding.fold_constants(model) + optimizer.remove_unused_nodes(model) + if_node = next(n for n in model.graph if n.op_type == "If") + then_branch = if_node.attributes["then_branch"].as_graph() + self.assertIn("temp", then_branch.initializers) + self.assertNotIn("temp", model.graph.initializers) + + # Make the condition constant (True) to trigger inlining of the then_branch. + const_true = ir.Value(name="const_true") + const_true.const_value = ir.Tensor(np.array(True)) + if_node.replace_input_with(0, const_true) + + # Second fold: the If is inlined; 'temp' must be moved to the main graph. + _constant_folding.fold_constants(model) + optimizer.remove_unused_nodes(model) + onnx.checker.check_model(ir.serde.serialize_model(model)) + self.assertIn("temp", model.graph.initializers) def test_fold_inside_if_branch(self): model = """