Repository navigation
[Fix][Relax] Drop integer-indexed axes in the Torch FX importer - #20991
Open
DiadkoShmek wants to merge 1 commit into
Open
DiadkoShmek wants to merge 1 commit into
DiadkoShmek wants to merge 1 commit into
Conversation
Treat integer subscripts as rank-reducing selections. Normalize negative positions and keep new axes in their PyTorch positions after the selections.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #20899.
The FX importer treats each integer in a tensor subscript as a length-one slice and keeps its axis. A bare integer also takes an early return that skips indexing. This can either reject valid Torch broadcasts (for example,
a[..., 0] + bfor shapes(3, 2, 4)and(3, 2)) or silently return the wrong rank when the operands remain broadcastable.Normalize negative integer positions against the axis extent and squeeze exactly the axes selected by integers. Compute the output positions of
Nonerelative to those removed axes. Slice objects continue to retain their axes.On
23a6dbe710dfcd445a090af1c49c8d7e8a57fb10, six new LLVM CPU tests throughtorch.compilefailed and one slice-only control passed. All seven pass with the patch and match Torch eager exactly. The full Dynamo and FX test modules went from 189 passed and 7 skipped to 196 passed and 7 skipped, including the seven added tests. Six exported-program indexing-related tests passed on both versions. A manual check also matched eager forx[..., -1]on two input shapes using one dynamically compiled graph.AI assistance (Codex) was used for the patch, tests, and verification.