Unintuitive behavior: Returned input aliases generate Identity nodes with generic names instead of preserving Python variable names

[patmjen]

Description

If a function returns a reassigned version of an input, the name of the output will get a generic name, i.e., return_val or return_val_[index] for multiple returns. This, to me, is unintuitive because it makes it harder to deduce the names of the generated outputs.

Consider this example case:

import onnx
from onnxscript import script, FLOAT
from onnxscript import opset20 as op

@script(default_opset=op)
def foo(x: FLOAT) -> FLOAT:
    y = x
    return y

onnx.checker.check_model(foo.to_model_proto())
print('Input:', [x.name for x in foo.to_model_proto().graph.input])
print('Output:', [x.name for x in foo.to_model_proto().graph.output])

This will print:

Input: ['x']
Output: ['return_val']

However, if I replace y = x with y = x + 0 or y = op.Identity(x), it will print Output: ['y']. I think it would be more intuitive if the generated output name was y since that matches the "intent" of the code. This will also match the behavior of using the y = op.Identity(x) which is what is "inserted" by the converter anyway.

Cause

The reason this happens comes from the internal ret function in Converter._translate_return_stmt:

onnxscript/onnxscript/converter.py

Lines 1070 to 1088 in 9dbf685

	def ret(exp, i, suffix):
	preferred_name = f"return_val{suffix}"
	return_var = self._translate_expr(exp, preferred_name).name
	val = self._lookup(return_var, self._source_of(exp), False)
	if val and val.kind == values.DynamicKind.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)
	for prev_output in self._current_fn.outputs:
	if prev_output.name == return_var:
	# ONNX does not allow duplicate output names.
	return_var = self._emit_copy(return_var, f"{return_var}_copy")
	break
	if self.returntype is None:
	t = None
	else:
	t = self.returntype[i]
	self.ir_builder.add_output(self._current_fn, return_var, t, self._source_of(stmt))
	return return_var

In the case of y = x, self._translate_expr(exp, preferred_name) will correctly deduce that y just refers to x so return_var will be 'x'. Therefore, val.kind will be an Input so a copy, i.e., an Identity op, will be emitted but it will use the generic preferred_name as the target.

Possible fix

An "easy" fix in this case could be to check if exp is an ast.Name and then use exp.id as the suggested name to emit_copy. This would result in an Identity op being generated with y as the output, so the model output is as expected (i.e., it matches the "intent" of the code). Since emit_copy still calls generate_unique_name, it should still be safe.

I fully realize this may not cover all cases and that it might be a personal opinion that this is unintuitive. Also, this fix suggestion does not cover the case of y = op.Identity(x); z = y; return z where the "intent" may be to have an output with name 'z'. So I would fully accept this not being changed.

Still, since there is an attempt at deducing the output names from the code (and an Identity op is inserted to make this a legal ONNX model/function/etc.), I thought this was worth bringing up. For now, i will just use explicit Identity operators. Thank you for the great library!

[justinchuby]

Thanks for noting this. Contributions welcomed!