[Relax][Frontend][TFLite] Add ATAN2 op and TFLite mapping#19485
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the atan2 operator with auto-broadcasting support across TVM's Relax and TOPI frameworks, including TFLite frontend integration, operator legalization, and comprehensive testing. The review feedback focuses on improving the documentation in broadcast.h, binary.py, and broadcast.py to explicitly clarify the parameter order as atan2(y, x), aligning with standard mathematical conventions to prevent potential misuse.
as4230
force-pushed
the
feat/relax-atan2-op
branch
from
28261c8 to
ac3f586
Compare
tlopex
left a comment
tlopex
left a comment
There was a problem hiding this comment.
Overalll LGTM! Please add legalize tests for atan2 in tests/python/relax/test_transform_legalize_ops_binary.py as well
This PR adds the ATAN2 operator to the Relax TFLite frontend.
Introduces relax.op.atan2 as a new binary elementwise primitive (TOPI
broadcast op, Relax registration, legalization to topi.atan2, script
parser support) and registers ATAN2 in the TFLite convert_map.
It reuses the TIR primitive tvm::atan2 so this PR is the higher-layer plumbing.
Validation:
python -m pytest tests/python/relax/test_op_binary.py
python -m pytest tests/python/relax/test_frontend_tflite.py -k binary
Addresses the ATAN2 item under apache#19412.
as4230
force-pushed
the
feat/relax-atan2-op
branch
from
ac3f586 to
b3d83c7
Compare
Done. I added test_atan2, test_atan2_symbolic, and test_atan2_primvalue in test_transform_legalize_ops_binary.py. |
2 participants
This PR adds the ATAN2 operator to the Relax TFLite frontend.
Introduces relax.op.atan2 as a new binary elementwise primitive (TOPI broadcast op, Relax registration, legalization to topi.atan2, script parser support) and registers ATAN2 in the TFLite convert_map. It reuses the TIR primitive tvm::atan2 so this PR is the higher-layer plumbing.
Validation:
python -m pytest tests/python/relax/test_op_binary.py
python -m pytest tests/python/relax/test_frontend_tflite.py -k binary
Addresses the ATAN2 item under #19412.