[Relax]feat: Implement FRelaxInferLayout for tile operator

[Dayuxiaoshui]

• Implement InferLayoutTile function to handle layout transformation for tile operator
• Use TransposeStrLike approach similar to repeat operator to correctly map repeats array
• Handle both same-dimension and dimension-expansion cases
• Add test case test_conv2d_tile to verify layout conversion from NCHW to NHWC
• Fixes the TODO at src/relax/op/tensor/manipulate.cc:1932

The implementation correctly transforms the repeats array when the input tensor's layout changes (e.g., from NCHW to NHWC), ensuring that repeat values are mapped to the correct dimensions in the new layout.

[gemini-code-assist]

Summary of Changes

Hello @Dayuxiaoshui, 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 significantly enhances the tile operator within Relax by introducing its FRelaxInferLayout capability. This ensures that when an input tensor's layout changes (e.g., from NCHW to NHWC), the tile operation's repeats array is intelligently remapped to align with the new layout. The change supports both scenarios where the tiling results in the same number of dimensions or expands them, leveraging a robust axis mapping strategy. This improvement is crucial for maintaining correctness and flexibility in tensor manipulations across different data layouts.

Highlights

• FRelaxInferLayout Implementation: Implemented the FRelaxInferLayout function for the tile operator, enabling it to correctly handle layout transformations.
• Repeats Array Transformation: The repeats array is now correctly transformed to match the new layout of the input tensor, using a TransposeStrLike approach similar to the repeat operator.
• Dimension Handling: The implementation supports both cases where the output dimensions are the same as the input and where dimensions are expanded due to tiling.
• Test Case Addition: A new test case, test_conv2d_tile, has been added to verify the tile operator's layout conversion, specifically from NCHW to NHWC.
• TODO Resolution: The long-standing TODO comment at src/relax/op/tensor/manipulate.cc:1932 regarding the implementation of FRelaxInferLayout for tile has been addressed and removed.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 implements layout inference for the tile operator, which is a great addition. The implementation correctly handles the tested case of same-dimension tiling. However, I found some issues in the implementation:

• The logic for handling len(repeats) < input.ndim is incorrect as it pads repeats at the end instead of the beginning as per the documentation.
• The logic for handling dimension expansion (len(repeats) > input.ndim) is incorrect.
• The implementation for same-dimension tiling is overly complex and hard to maintain.

I've provided a simplified and corrected implementation for InferLayoutTile that addresses these issues. I've also suggested adding more test cases to cover the scenarios where the current implementation would fail, ensuring the feature is robust.

[Dayuxiaoshui]

@tlopex