From 01c46b5954371fc7f20d8fe28bdccd55ef233279 Mon Sep 17 00:00:00 2001 From: TANGBUDU <79953480+TANGBUDU@users.noreply.github.com> Date: Fri, 28 Aug 2026 18:26:37 +0200 Subject: [PATCH 1/2] Preserve returned input alias names --- onnxscript/_internal/converter.py | 7 ++++++- onnxscript/_internal/converter_test.py | 17 +++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/onnxscript/_internal/converter.py b/onnxscript/_internal/converter.py index c2f6b0fb63..3789bb5f47 100644 --- a/onnxscript/_internal/converter.py +++ b/onnxscript/_internal/converter.py @@ -1108,7 +1108,12 @@ def ret(exp, i, suffix): if val.value.is_graph_input(): # In ONNX, a graph-input cannot be an output of the graph. # We need to insert a copy. - return_var = self._emit_copy(return_var, preferred_name) + copy_name = ( + exp.id + if isinstance(exp, ast.Name) and exp.id != return_var.name + else preferred_name + ) + return_var = self._emit_copy(return_var, copy_name) for prev_output in self._current_fn.outputs: if prev_output.name == return_var.name: # ONNX does not allow duplicate output names. diff --git a/onnxscript/_internal/converter_test.py b/onnxscript/_internal/converter_test.py index c1fe276ef5..ea5eb1f339 100644 --- a/onnxscript/_internal/converter_test.py +++ b/onnxscript/_internal/converter_test.py @@ -607,6 +607,23 @@ def duplicate_output(X): outputs = duplicate_output.to_function_proto().output self.assertNotEqual(outputs[0], outputs[1]) + def test_returned_input_alias_preserves_name(self): + @script(default_opset=op) + def returned_alias(X): + Y = X + return Y + + function_proto = returned_alias.to_function_proto() + self.assertEqual(function_proto.output[0], "Y") + self.assertEqual(function_proto.node[-1].op_type, "Identity") + self.assertEqual(function_proto.node[-1].output[0], "Y") + + @script(default_opset=op) + def returned_input(X): + return X + + self.assertEqual(returned_input.to_function_proto().output[0], "return_val") + def test_bool_attr_promotion(self): @script() def if_then_else(flag: bool, Y, Z): From cbcda16aa93281d0c5e5dcb50baa96951e3a2f08 Mon Sep 17 00:00:00 2001 From: TANGBUDU <79953480+TANGBUDU@users.noreply.github.com> Date: Mon, 7 Sep 2026 10:23:27 +0200 Subject: [PATCH 2/2] Preserve returned aliases after rebinding the input name Check the resolved ONNX value rather than looking its name up in the mutable Python symbol table. Cover input rebinding and collisions with existing input names, including duplicate returns and numerical parity. Assisted-by: ChatGPT --- onnxscript/_internal/converter.py | 19 ++++++-------- onnxscript/_internal/converter_test.py | 35 ++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 11 deletions(-) diff --git a/onnxscript/_internal/converter.py b/onnxscript/_internal/converter.py index 3789bb5f47..f584315a5f 100644 --- a/onnxscript/_internal/converter.py +++ b/onnxscript/_internal/converter.py @@ -1103,17 +1103,14 @@ def check_num_outputs(n): def ret(exp, i, suffix): preferred_name = f"return_val{suffix}" return_var = self._translate_expr(exp, preferred_name) - val = self._lookup(return_var.name, self._source_of(exp), raise_exception=False) - if isinstance(val, values.SymbolValue) and isinstance(val.value, ir.Value): - if val.value.is_graph_input(): - # In ONNX, a graph-input cannot be an output of the graph. - # We need to insert a copy. - copy_name = ( - exp.id - if isinstance(exp, ast.Name) and exp.id != return_var.name - else preferred_name - ) - return_var = self._emit_copy(return_var, copy_name) + if return_var.is_graph_input(): + # Use the resolved ONNX value: the Python input name may have been rebound. + copy_name = ( + exp.id + if isinstance(exp, ast.Name) and exp.id != return_var.name + else preferred_name + ) + return_var = self._emit_copy(return_var, copy_name) for prev_output in self._current_fn.outputs: if prev_output.name == return_var.name: # ONNX does not allow duplicate output names. diff --git a/onnxscript/_internal/converter_test.py b/onnxscript/_internal/converter_test.py index ea5eb1f339..fba141b428 100644 --- a/onnxscript/_internal/converter_test.py +++ b/onnxscript/_internal/converter_test.py @@ -624,6 +624,41 @@ def returned_input(X): self.assertEqual(returned_input.to_function_proto().output[0], "return_val") + def test_returned_input_alias_after_rebinding_input(self): + @script(default_opset=op) + def returned_alias(X: FLOAT[2]) -> FLOAT[2]: + Y = X + X = op.Neg(X) + return Y + + model = returned_alias.to_model_proto() + onnx.checker.check_model(model, full_check=True) + self.assertEqual(model.graph.output[0].name, "Y") + self.assertEqual(model.graph.node[-1].op_type, "Identity") + self.assertEqual(list(model.graph.node[-1].input), ["X"]) + x = np.array([1.0, -2.0], dtype=np.float32) + actual = create_cpu_inference_session(model.SerializeToString()).run(None, {"X": x}) + np.testing.assert_array_equal(actual[0], x) + + def test_returned_input_alias_name_collision(self): + @script(default_opset=op) + def returned_alias(X: FLOAT[2], Y: FLOAT[2]) -> (FLOAT[2], FLOAT[2]): + Y = X + return Y, Y + + model = returned_alias.to_model_proto() + onnx.checker.check_model(model, full_check=True) + outputs = [value.name for value in model.graph.output] + self.assertEqual(len(set(outputs)), 2) + self.assertTrue(all(name.startswith("Y_") for name in outputs)) + self.assertTrue(set(outputs).isdisjoint({"X", "Y"})) + x = np.array([1.0, -2.0], dtype=np.float32) + actual = create_cpu_inference_session(model.SerializeToString()).run( + None, {"X": x, "Y": -x} + ) + for output in actual: + np.testing.assert_array_equal(output, x) + def test_bool_attr_promotion(self): @script() def if_then_else(flag: bool, Y, Z):