Fix aten.amax lowering issue

[cccclai]

Summary:
There was an error when lowering amax around this line input_tensor = self.get_tensor(input_node, node) and the issue is that we're trying to permute the tensor inside node_visitors, op_node.meta[QCOM_AXIS_ORDER] is (0, 1), however, tensor.shape is (1, 980, 49).

Rollback Plan:

Differential Revision: D80187368

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/13381

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit c775e98 with merge base 45e7810 ():

NEW FAILURES - The following jobs have failed:

• Build documentation / build (buck2) / Build doc (gh)
  At least one of the pre-conditions you specified did not hold
• pull / unittest-arm-backend-with-no-fvp (test_pytest_ops) / linux-job (gh)
  RuntimeError: Command docker exec -t 35b7396f1bf047158ffcdfad0c168c96dd2863b4ae488602a467c2cf731fb6da /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D80187368

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[cccclai]

This PR fix the lowering issue, however, not quite sure a good unit test for it. Try to use the following code but looks like passing without this PR. Do you know what might trigger the error?

class ConvAMax2D(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, dim=None, keepdim=False):
        super().__init__()
        self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size, padding=kernel_size//2)  
        self.dim = dim
        self.keepdim = keepdim

    def forward(self, x):
        x = self.conv(x) 
        return torch.amax(x, dim=self.dim, keepdim=self.keepdim)

sample_input = torch.randn(2, 3, 64, 64) 
module = ConvAMax2D(in_channels=3, out_channels=16, kernel_size=3, dim=1, keepdim=False)
output = module(sample_input)

Edit:
Actually it does trigger the error, add it in the unit test

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D80187368

[cccclai]

Can I get a review on this PR? @haowhsu-quic @shewu-quic @winskuo-quic @DannyYuyang-quic

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D80187368

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D80187368