From 0aa6f15038a724e9e239ebcf2f00e0b4df8ce8f4 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 25 Feb 2026 06:58:50 -0800 Subject: [PATCH 1/4] Add input() and add_output() methods to GraphBuilder Signed-off-by: Justin Chu --- onnxscript/_internal/builder.py | 46 +++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/onnxscript/_internal/builder.py b/onnxscript/_internal/builder.py index fb81f3451f..a4cd6dd98a 100644 --- a/onnxscript/_internal/builder.py +++ b/onnxscript/_internal/builder.py @@ -152,6 +152,52 @@ def initializer( self._graph.register_initializer(value) return value + def input( + self, + name: str, + dtype: ir.DataType | None = None, + shape: ir.Shape | Sequence[int | str | None] | None = None, + *, + type: ir.TypeProtocol | None = None, + const_value: ir.TensorProtocol | None = None, + metadata_props: dict[str, str] | None = None, + ) -> ir.Value: + """Create an input to the graph and return the corresponding ir.Value. + + Args: + name: The name of the value. + dtype: The data type of the TensorType of the value. This is used only when type is None. + shape: The shape of the value. + type: The type of the value. Only one of dtype and type can be specified. + const_value: The constant tensor that initializes the value. Supply this argument + when you want to create an initializer. The type and shape can be obtained from the tensor. + metadata_props: The metadata properties that will be serialized to the ONNX proto. + + Returns: + A Value object. + """ + value = ir.val( + name=name, + dtype=dtype, + shape=shape, + type=type, + const_value=const_value, + metadata_props=metadata_props, + ) + self._graph.inputs.append(value) + return value + + def add_output(self, value: ir.Value, name: str | None) -> None: + """Add an output to the graph. + + Args: + value: The ir.Value to add as an output. + name: The name to assign to the output value. If None, no renaming is done. + """ + if name: + value.name = name + self._graph.outputs.append(value) + def _input_to_ir_value( self, value: VALUE_LIKE, like_type: ir.Value | None = None ) -> ir.Value: From 2ee8a41f4ea6f0a00d44ab89a9a33b4292c02c09 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 25 Feb 2026 07:02:52 -0800 Subject: [PATCH 2/4] Add tests for GraphBuilder input/output/initializer APIs --- onnxscript/_internal/builder_test.py | 59 ++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/onnxscript/_internal/builder_test.py b/onnxscript/_internal/builder_test.py index 1ff8460b67..63f262356b 100644 --- a/onnxscript/_internal/builder_test.py +++ b/onnxscript/_internal/builder_test.py @@ -541,6 +541,65 @@ def test_output_names_are_unique_for_same_op_type(self): names = [t1.name, t2.name, t3.name] self.assertEqual(len(set(names)), 3) + def test_input_creates_and_registers_graph_input(self): + """Test that GraphBuilder.input creates and appends a graph input value.""" + graph = ir.Graph( + name="test_model", + inputs=[], + outputs=[], + nodes=[], + opset_imports={"": _default_opset_version}, + ) + graph_builder = builder.GraphBuilder(graph) + + value = graph_builder.input("data", dtype=ir.DataType.FLOAT, shape=[2, 3]) + + self.assertEqual(value.name, "data") + self.assertEqual(value.type.dtype, ir.DataType.FLOAT) + self.assertEqual(list(value.shape), [2, 3]) + self.assertEqual(len(graph.inputs), 1) + self.assertIs(graph.inputs[0], value) + + def test_add_output_renames_and_registers_output(self): + """Test that GraphBuilder.add_output renames (optionally) and appends outputs.""" + graph = ir.Graph( + name="test_model", + inputs=[], + outputs=[], + nodes=[], + opset_imports={"": _default_opset_version}, + ) + graph_builder = builder.GraphBuilder(graph) + + output = ir.Value(name="old_name") + graph_builder.add_output(output, "new_name") + + self.assertEqual(output.name, "new_name") + self.assertEqual(len(graph.outputs), 1) + self.assertIs(graph.outputs[0], output) + + def test_initializer_qualification_behavior(self): + """Test that GraphBuilder.initializer qualifies names unless explicitly disabled.""" + graph = ir.Graph( + name="test_model", + inputs=[], + outputs=[], + nodes=[], + opset_imports={"": _default_opset_version}, + ) + graph_builder = builder.GraphBuilder(graph) + + graph_builder.push_module("layer1") + qualified = graph_builder.initializer(ir.tensor([1.0], name="w"), name="weight") + unqualified = graph_builder.initializer( + ir.tensor([2.0], name="b"), name="bias", qualify=False + ) + + self.assertEqual(qualified.name, "layer1.weight") + self.assertEqual(unqualified.name, "bias") + self.assertIn("layer1.weight", graph.initializers) + self.assertIn("bias", graph.initializers) + def test_multi_output_names_are_unique(self): """Test that multi-output ops produce unique names with counter suffix.""" op, x, y = _create_builder_with_inputs() From af2a1abe0b1b1e3c787cd35e39be0f6dd0819559 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 25 Feb 2026 11:21:04 -0800 Subject: [PATCH 3/4] Register initializer Signed-off-by: Justin Chu --- onnxscript/_internal/builder.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/onnxscript/_internal/builder.py b/onnxscript/_internal/builder.py index a4cd6dd98a..19317bbfe1 100644 --- a/onnxscript/_internal/builder.py +++ b/onnxscript/_internal/builder.py @@ -185,6 +185,8 @@ def input( metadata_props=metadata_props, ) self._graph.inputs.append(value) + if const_value is not None: + self._graph.register_initializer(value) return value def add_output(self, value: ir.Value, name: str | None) -> None: From e725875cca5570fce3b5603b545178e568ecef6c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 25 Feb 2026 11:23:25 -0800 Subject: [PATCH 4/4] test Signed-off-by: Justin Chu --- onnxscript/_internal/builder_test.py | 37 ++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/onnxscript/_internal/builder_test.py b/onnxscript/_internal/builder_test.py index 63f262356b..790bfc6f65 100644 --- a/onnxscript/_internal/builder_test.py +++ b/onnxscript/_internal/builder_test.py @@ -560,6 +560,43 @@ def test_input_creates_and_registers_graph_input(self): self.assertEqual(len(graph.inputs), 1) self.assertIs(graph.inputs[0], value) + def test_input_with_const_value_registers_initializer(self): + """Test that GraphBuilder.input registers initializer when const_value is provided.""" + graph = ir.Graph( + name="test_model", + inputs=[], + outputs=[], + nodes=[], + opset_imports={"": _default_opset_version}, + ) + graph_builder = builder.GraphBuilder(graph) + + const_tensor = ir.tensor([1.0, 2.0], dtype=ir.DataType.FLOAT, name="const_data") + value = graph_builder.input("const_input", const_value=const_tensor) + + self.assertEqual(len(graph.inputs), 1) + self.assertIs(graph.inputs[0], value) + self.assertIn("const_input", graph.initializers) + self.assertIs(graph.initializers["const_input"], value) + self.assertIs(value.const_value, const_tensor) + + def test_input_without_const_value_does_not_register_initializer(self): + """Test that GraphBuilder.input does not register initializer without const_value.""" + graph = ir.Graph( + name="test_model", + inputs=[], + outputs=[], + nodes=[], + opset_imports={"": _default_opset_version}, + ) + graph_builder = builder.GraphBuilder(graph) + + value = graph_builder.input("regular_input", dtype=ir.DataType.FLOAT, shape=[2]) + + self.assertEqual(len(graph.inputs), 1) + self.assertIs(graph.inputs[0], value) + self.assertNotIn("regular_input", graph.initializers) + def test_add_output_renames_and_registers_output(self): """Test that GraphBuilder.add_output renames (optionally) and appends outputs.""" graph = ir.Graph(