Skip to content

Fix aten.copy: wrap CastLike with Identity to prevent IR value aliasing#2875

Open
Copilot wants to merge 2 commits intomainfrom
copilot/fix-torch-export-valueerror
Open

Fix aten.copy: wrap CastLike with Identity to prevent IR value aliasing#2875
Copilot wants to merge 2 commits intomainfrom
copilot/fix-torch-export-valueerror

Conversation

Copy link
Copy Markdown
Contributor

Copilot AI commented Mar 31, 2026

torch.onnx.export raised ValueError when exporting an ExportedProgram containing register_buffer + in-place copy (copy_), because aten_copy's op.CastLike(src, self) returns the same IR value object as src when both tensors share the same dtype — no new ONNX node is created. A subsequent unconditional rename of that value clobbered the original buffer placeholder name (e.g. b_prompt_featcopy), causing initializer registration to fail.

Change

  • onnxscript/function_libs/torch_lib/ops/core.py: Wrap the CastLike result in op.Identity() so aten_copy always produces a fresh IR value, leaving the source value's name intact.
# Before
def aten_copy(self: TTensor, src: TTensor2, non_blocking: bool = False) -> TTensor:
    return op.CastLike(src, self)

# After
def aten_copy(self: TTensor, src: TTensor2, non_blocking: bool = False) -> TTensor:
    return op.Identity(op.CastLike(src, self))

This matches the existing pattern in aten__to_copy, which already returns op.Identity(self) for the same-dtype case.

Copilot AI changed the title [WIP] Fix ValueError on torch.onnx.export for ExportedProgram Fix aten.copy: wrap CastLike with Identity to prevent IR value aliasing Mar 31, 2026
Copilot AI requested a review from justinchuby March 31, 2026 16:10
@justinchuby justinchuby reopened this Mar 31, 2026
@justinchuby justinchuby marked this pull request as ready for review March 31, 2026 16:27
@codecov
Copy link
Copy Markdown

codecov bot commented Mar 31, 2026

Codecov Report

❌ Patch coverage is 0% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 71.96%. Comparing base (1077da7) to head (83ee953).
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
onnxscript/function_libs/torch_lib/ops/core.py 0.00% 1 Missing ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##             main    #2875   +/-   ##
=======================================
  Coverage   71.96%   71.96%           
=======================================
  Files         239      239           
  Lines       29224    29224           
  Branches     2878     2878           
=======================================
  Hits        21031    21031           
  Misses       7216     7216           
  Partials      977      977           

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

@justinchuby
Copy link
Copy Markdown
Collaborator

Actually we should probably fix the castlike evaluation logic to emit an identity also.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

[ONNX] ExportedProgram with register_buffer + in-place copy fails on torch.onnx.export

2 participants