feat: support amax dynamo converter

[zewenli98]

Description

Support amax dynamo converter.

• Function Schema:
  torch.ops.aten.amax.default

• Original PyTorch API:
  https://pytorch.org/docs/stable/generated/torch.amax.html

Fixes issue #2095

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[narendasan]

@gs-olive can this check be done in a validator?

[gs-olive]

This check can be done in a validator, but in this context it would make the most sense for it to be done here, since we can support cases where this argument is both present and absent.

[gs-olive]

This check can be done in a validator, but in this context it would make the most sense for it to be done here, since we can support cases where this argument is both present and absent.

[gs-olive]

Converter looks great, and the testing is very robust! Made a few small comments

[gs-olive]

Is this required for the IReduceLayer?

[zewenli98]

Yes, it's required because TypeError: Provided an unsupported data type as an input data type (support: bool, int32, long, half, float), got: torch.int8 if I test int8. Even though the doc says Reduce Layer supports int32, but I tested it doesn't support.

[zewenli98]

Converter looks great, and the testing is very robust! Made a few small comments

Thanks for your review! I've solved all the issues and pushed.

[zewenli98]

It looks like failed on C++ tests that's nothing to do with this commit?

[gs-olive]

I reran that C++ test here

[zewenli98]

Passed! please let me know if anything need to fix in this PR.

[gs-olive]

Based on this schema, it seems that the dimension (dim) can also be not present. For instance, the following model creates the subsequent graph:

class argmax(torch.nn.Module):
   def forward(self, x):
        return torch.argmax(x)

"""
graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %argmax : [num_users=1] = call_function[target=torch.ops.aten.argmax.default](args = (%arg0_1,), kwargs = {})
    return (argmax,)
"""

If we cannot support this case, you can add a capability_validator function to this decorator, which will note that case as unsupported. Roughly, that could be something like:

def amax_param_validator(amax_node: Node) -> bool:
    return len(amax_node.args) >= 2

@dynamo_tensorrt_converter(
    torch.ops.aten.amax.default, capability_validator=amax_param_validator
)  # type: ignore[misc]

[zewenli98]

@gs-olive Thank you so much George! That's a good catch. I fixed and learned a lot from these details recently 😃

[gs-olive]

No problem - glad to hear it!

[gs-olive]

Looks great!