support argmax converter

[bowang007]

Description

Support argmax converter

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

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/argmax.py	2023-09-05 22:31:02.244529+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/argmax.py	2023-09-05 22:33:23.441716+00:00
@@ -23,18 +23,15 @@
    dim: int = 0,
    keep_dim: bool = False,
) -> TRTTensor:
    if not isinstance(input, TRTTensor):
        raise RuntimeError(
-            f"argmax received input {input} that is not part "
-            "of the TensorRT region!"
+            f"argmax received input {input} that is not part " "of the TensorRT region!"
        )
    if dim < 0:
        dim = len(tuple(input.shape)) + dim
    reduce_mask = 1 << dim
    topk_layer = network.add_topk(input, trt.TopKOperation.MAX, 1, reduce_mask)

    set_layer_name(topk_layer, target, name)

    return topk_layer.get_output(1)
-    
-    
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/converters/test_argmax.py	2023-09-05 22:31:02.264529+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/converters/test_argmax.py	2023-09-05 22:33:26.764451+00:00
@@ -2,33 +2,23 @@
import torch.nn as nn
from parameterized import parameterized
from torch.testing._internal.common_utils import run_tests
from harness import DispatchTestCase

+
class TestArgmaxConverter(DispatchTestCase):
-    @parameterized.expand(
-            [
-                ("dim_0_keep_dim_false", (3, 4), 0, False)
-            ]
-    )
-
+    @parameterized.expand([("dim_0_keep_dim_false", (3, 4), 0, False)])
    def test_argmax(self, _, input_shape, dim, keep_dim):
        class ArgMax(nn.Module):
            def __init__(self):
                super().__init__()

-            def forward(self, input): 
+            def forward(self, input):
                return torch.argmax(input, dim, keep_dim)
-            

        input = [torch.randn(*input_shape)]

-        self.run_test(
-            ArgMax(),
-            input, 
-            expected_ops={torch.ops.aten.argmax.default}
-        )
+        self.run_test(ArgMax(), input, expected_ops={torch.ops.aten.argmax.default})
+

if __name__ == "__main__":
-    run_tests()  
-
-
+    run_tests()

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[gs-olive]

Switch to dim: Optional[int] = None since this is the default dim, as per the documentation. Alternatively, if this converter cannot support reducing over all dimensions, you can add a capability_validator to the converter to disallow inputs where the dim is not specified or non-integral.

[bowang007]

I used dim: Union[int, None], is that ok?

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[bowang007]

Hey @gs-olive I will be OOO next week.
I think this update covers all edge cases.
Please feel free to merge if this is good to go. Thanks!

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[github-actions]

Code conforms to Python style guidelines

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[gs-olive]

Looks good to me

[zewenli98]

Looks good overall, but I left some comments about using our new APIs and small fixes.

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[github-actions]

Code conforms to Python style guidelines

[github-actions]

Code conforms to C++ style guidelines

[zewenli98]

Thanks for the fix! I found a small bug here. Other looks good to me!

[github-actions]

Code conforms to C++ style guidelines

[github-actions]

Code conforms to Python style guidelines

[zewenli98]

LGTM!