[Relax][PyTorch] Add support for decomposed operators and fix IR of ops tests(5)

[tlopex]

This pr fixes ops like tile, transpose and so on.

[gemini-code-assist]

Summary of Changes

Hello @tlopex, I'm Gemini Code Assist¹! 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 Relax. It updates the expected Intermediate Representation (IR) in several test cases to accurately reflect how operators like stack, tile, fill, masked_fill, and type conversions are broken down into more primitive operations. This ensures proper handling and verification of these decomposed operators within the Relax framework.

Highlights

• Operator Decomposition Support: Enabled run_ep_decomposition for numerous verify_model calls across various tests, indicating the activation of operator decomposition for a wide range of operations.
• IR Adjustments for stack Operator: The expected Intermediate Representation (IR) for R.stack operations was updated to reflect their decomposition into R.concat and R.reshape primitives.
• tile Operator Argument Fix: The R.tile operator's repeats argument is now explicitly named in the expected IR, improving clarity and correctness in its representation.
• fill and masked_fill IR Changes: R.full was replaced with R.full_like for fill operations, and masked_fill now uses R.const for the fill value in its expected IR.
• Simplified Type Conversion IR: Explicit R.astype calls were removed from the expected IR for various type conversion tests, suggesting a more streamlined representation or implicit handling.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request updates tests for the PyTorch frontend to support decomposed operators by enabling run_ep_decomposition=True. The changes involve updating the expected Relax IR for several operators like stack, tile, fill, masked_fill, and type casting ops to match the output of PyTorch's decomposition pass. Most changes correctly reflect more efficient or accurate decompositions. However, I've pointed out an inconsistency in the decomposition of torch.stack for different axes and suggested a more canonical approach for better robustness and consistency.

[gemini-code-assist]

The decomposition of torch.stack with axis=0 into concat and reshape is valid for contiguous tensors, but it's not as canonical as using expand_dims and concat. The decomposition for axis=-1 in Expected3 uses expand_dims and concat, which is a more direct translation of the stack operation's semantics of inserting a new dimension. For consistency and clarity, it would be better to use the same expand_dims and concat approach for axis=0. This would make the decomposition logic more robust and easier to understand across different axes.

For example:

lv: R.Tensor((1, 2, 3), dtype="float32") = R.expand_dims(x, axis=0)
lv1: R.Tensor((1, 2, 3), dtype="float32") = R.expand_dims(y, axis=0)
lv2: R.Tensor((2, 2, 3), dtype="float32") = R.concat((lv, lv1), axis=0)
gv: R.Tuple(R.Tensor((2, 2, 3), dtype="float32")) = (lv2,)

[gemini-code-assist]

Similar to the axis=0 case, the decomposition of torch.stack with axis=1 into concat and reshape is less canonical than using expand_dims and concat. Using a consistent decomposition strategy across all axes would improve the robustness and readability of the translated IR.

For example:

lv: R.Tensor((2, 1, 3), dtype="float32") = R.expand_dims(x, axis=1)
lv1: R.Tensor((2, 1, 3), dtype="float32") = R.expand_dims(y, axis=1)
lv2: R.Tensor((2, 2, 3), dtype="float32") = R.concat((lv, lv1), axis=1)
gv: R.Tuple(R.Tensor((2, 2, 3), dtype="float32")) = (lv2,)