Add OutputFixPass to fix invalid graph outputs

[justinchuby]

This pull request introduces a new pass, OutputFixPass, to the ONNX IR transformation pipeline. The main purpose of this pass is to ensure that the model's outputs conform to ONNX specifications by inserting Identity nodes where necessary. This helps to fix cases where graph inputs are directly used as outputs or where the same value is used multiple times as an output, both of which are not allowed by ONNX.

The most important changes include:

New Pass Implementation

• Added a new file, output_fix.py, implementing OutputFixPass. This pass automatically inserts Identity nodes in two scenarios: when a graph input is directly used as an output, and when a value is used multiple times as an output. The pass applies to both the main graph and all subgraphs/functions.

AI use

I used copilot to create the tests and some of the logic.

[Copilot]

Pull request overview

This PR introduces OutputFixPass, a new transformation pass that ensures ONNX model validity by inserting Identity nodes when graph inputs are directly used as graph outputs. This addresses an ONNX specification requirement where direct input-to-output connections are not permitted.

Key Changes:

• Implements OutputFixPass to detect and fix direct input-to-output connections by inserting Identity nodes
• Handles main graphs, subgraphs, and function bodies recursively
• Preserves output metadata (name, shape, type, doc_string, and metadata_props) when inserting Identity nodes

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

File	Description
src/onnx_ir/passes/common/output_fix.py	Implements the new OutputFixPass with logic to insert Identity nodes between direct input-output connections
src/onnx_ir/passes/common/output_fix_test.py	Comprehensive test suite covering various scenarios including subgraphs, functions, and edge cases
src/onnx_ir/passes/common/__init__.py	Registers and exports OutputFixPass in the common passes module

[codecov]

Codecov Report

❌ Patch coverage is 93.75000% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.68%. Comparing base (ac8138e) to head (791cdec).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
src/onnx_ir/passes/common/output_fix.py	93.54%	1 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #269      +/-   ##
==========================================
+ Coverage   77.33%   77.68%   +0.35%     
==========================================
  Files          40       41       +1     
  Lines        5042     5105      +63     
  Branches     1009     1022      +13     
==========================================
+ Hits         3899     3966      +67     
+ Misses        853      844       -9     
- Partials      290      295       +5     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

[Copilot]

The condition if output_usages[output] > 1: is redundant here. Since we're inside the loop iterating through outputs and we've already seen this output once (it's not in the seen set check), we already know that output_usages[output] > 1. This check will always be true at this point in the code flow.

Suggested change

	if output_usages[output] > 1:
	# Create an Identity node
	identity_node = ir.node("Identity", inputs=[output])
	identity_output = identity_node.outputs[0]
	# Copy metadata from the original output
	# TODO: Use a better unique naming strategy if needed
	identity_output.name = f"{output.name}_alias_{i}"
	identity_output.shape = output.shape
	identity_output.type = output.type
	identity_output.metadata_props.update(output.metadata_props)
	identity_output.doc_string = output.doc_string
	# Add the node to the graph
	graph.append(identity_node)
	graph.outputs[i] = identity_output
	logger.debug(
	"Added Identity node for graph output '%s' used multiple times", output
	)
	modified = True
	# Create an Identity node
	identity_node = ir.node("Identity", inputs=[output])
	identity_output = identity_node.outputs[0]
	# Copy metadata from the original output
	# TODO: Use a better unique naming strategy if needed
	identity_output.name = f"{output.name}_alias_{i}"
	identity_output.shape = output.shape
	identity_output.type = output.type
	identity_output.metadata_props.update(output.metadata_props)
	identity_output.doc_string = output.doc_string
	# Add the node to the graph
	graph.append(identity_node)
	graph.outputs[i] = identity_output
	logger.debug(
	"Added Identity node for graph output '%s' used multiple times", output
	)
	modified = True

[Copilot]

Potential name collision issue: Using f"{output.name}_alias_{i}" where i is the output index could cause collisions if there's already a value with a similar name pattern in the graph. Consider using a more robust unique naming strategy that checks for existing names in the graph.

Suggested change

	# TODO: Use a better unique naming strategy if needed
	identity_output.name = f"{output.name}_alias_{i}"
	# Use a robust unique naming strategy
	identity_output.name = _get_unique_value_name(graph, f"{output.name}_alias_{i}")