[Relax][PyTorch] Add support for decomposed operators in extended unary ops tests#18400
Conversation
Summary of ChangesHello @tlopex, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request integrates support for decomposed operators within the PyTorch frontend for TVM Relax. By adding specific mappings and decomposition logic for several unary operations, it improves the accuracy and flexibility of translating PyTorch models. The accompanying test updates ensure that these decomposed operations are correctly handled and verified, contributing to a more robust and comprehensive frontend. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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 PR adds support for several decomposed PyTorch operators in the Relax frontend and updates the test suite to verify their translation, particularly when run_ep_decomposition is enabled. The changes are generally good and increase coverage for decomposed ops. However, I've found a few issues in the test expectations: one of the dropout tests seems to have incorrect logic that always produces a zero tensor, and the expected output for inplace operations is a tuple of two tensors instead of one, which is confusing. There's also a minor inconsistency in the relu6 tests. Please see my detailed comments.
| class expected_dropout_for_3: | ||
| @R.function | ||
| def main( | ||
| input: R.Tensor((1, 3, 10, 10), dtype="float32") | ||
| ) -> R.Tuple( | ||
| R.Tensor((1, 3, 10, 10), dtype="float32"), R.Tensor((1, 3, 10, 10), dtype="float32") | ||
| ): | ||
| # block 0 | ||
| with R.dataflow(): | ||
| lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.zeros( | ||
| R.shape([1, 3, 10, 10]), dtype="float32" | ||
| ) | ||
| lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide( | ||
| lv, R.const(0.5, "float32") | ||
| ) | ||
| lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(input, lv1) | ||
| gv: R.Tuple( | ||
| R.Tensor((1, 3, 10, 10), dtype="float32"), | ||
| R.Tensor((1, 3, 10, 10), dtype="float32"), | ||
| ) = (lv2, lv2) | ||
| R.output(gv) | ||
| return gv |
There was a problem hiding this comment.
The expected IR for Dropout3 seems incorrect as it will always produce a zero tensor. The logic R.zeros(...) followed by R.divide and R.multiply results in a zero tensor, which does not correctly model the behavior of dropout, even for deterministic testing. This seems to be caused by the translation of torch.empty_like to relax.op.zeros_like, which then leads to a zero mask after the bernoulli operation. This testing logic for dropout should be revisited to correctly reflect its behavior.
| class expected_hardswish_for_3: | ||
| @R.function | ||
| def main( | ||
| input: R.Tensor((1, 3, 10, 10), dtype="float32") | ||
| ) -> R.Tuple( | ||
| R.Tensor((1, 3, 10, 10), dtype="float32"), R.Tensor((1, 3, 10, 10), dtype="float32") | ||
| ): | ||
| with R.dataflow(): | ||
| lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add( | ||
| input, R.const(3.0, "float32") | ||
| ) | ||
| lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip( | ||
| lv, R.prim_value(0), R.prim_value(T.float64("inf")) | ||
| ) | ||
| lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip( | ||
| lv1, R.prim_value(T.float64("-inf")), R.prim_value(6) | ||
| ) | ||
| lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(input, lv2) | ||
| lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide( | ||
| lv3, R.const(6.0, "float32") | ||
| ) | ||
| gv: R.Tuple( | ||
| R.Tensor((1, 3, 10, 10), dtype="float32"), | ||
| R.Tensor((1, 3, 10, 10), dtype="float32"), | ||
| ) = (lv4, lv4) | ||
| R.output(gv) | ||
| return gv |
There was a problem hiding this comment.
The expected IR for the inplace hardswish_ operation returns a tuple of two identical tensors. This is inconsistent with the original PyTorch model Hardswish3, which returns a single tensor. This pattern also appears in the test for relu6_ (expected_relu6_3). While this might be an intended consequence of how inplace operations are handled after decomposition, it's confusing. The translated function should ideally match the return signature of the original model. Could you clarify if this is intended, or adjust the expected IR to return a single tensor?
| return gv | ||
|
|
||
| verify_model(ReLU6_1(), example_args, {}, expected_relu6_1, run_ep_decomposition=True) | ||
| verify_model(ReLU6_2(), example_args, {}, expected_relu6_2, run_ep_decomposition=True) |
There was a problem hiding this comment.
There seems to be an inconsistency in how relu6 is handled. The test for ReLU6_1 (using nn.ReLU6) shows that it's decomposed into relax.op.clip, which is expected with run_ep_decomposition=True. However, this test for ReLU6_2 (using F.relu6) is expected to be translated to relax.op.nn.relu6, indicating it's not decomposed. Since both tests use run_ep_decomposition=True, this difference is surprising. Is this discrepancy in PyTorch's decomposition behavior expected?
|
cc @mshr-h |
2 participants
This PR adds support for decomposed operators that ops needed in the extended unary ops tests.