[Relax][PyTorch] Add atan2 converter#19850
Open
javierdejesusda wants to merge 1 commit into
Open
Conversation
torch.atan2 was not registered in either the ExportedProgram or FX frontend, so importing a model that uses it failed with an "Unsupported function types" error. The relax.op.atan2 operator already exists and legalizes to topi.atan2, so the frontends only needed to route the op to it. Register atan2 in the FX frontend and atan2.default in the ExportedProgram frontend, reusing the shared _binary_op helper (the same pattern as the existing maximum/minimum/logaddexp converters), and add a structural test in both test_frontend_from_fx.py and test_frontend_from_exported_program.py.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request adds support for the atan2 operator in both the FX translator and the exported program translator within the TVM Relax PyTorch frontend. It also includes corresponding unit tests to verify the translation. No review comments were provided, and I have no additional feedback to offer as the changes are clean and well-tested.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
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
1 participant
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.
Motivation
torch.atan2was not registered in either the ExportedProgram or FX frontend,so importing a model that uses it failed with an "Unsupported function types"
error. The
relax.op.atan2operator already exists and legalizes totopi.atan2, so the frontends only needed to route the op to it.Changes
atan2in the FX frontend andatan2.defaultin the ExportedProgramfrontend, reusing the shared
_binary_ophelper (the same pattern as theexisting
maximum/minimum/logaddexpconverters).test_frontend_from_fx.pyandtest_frontend_from_exported_program.py.