fix dim bugs and update tests

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -645,7 +645,7 @@ def aten_ops_amax(
         SourceIR.ATEN,
         name,
         args[0],
-        args[1],
+        args_bounds_check(args, 1, replacement=[]),
         args_bounds_check(args, 2, replacement=False),
     )
 

py/torch_tensorrt/dynamo/conversion/impl/reduce.py

@@ -27,6 +27,9 @@ def amax(
     ):
         input_val = cast_trt_tensor(ctx, input_val, trt.float32, name)
 
+    if dim is None or (isinstance(dim, (tuple, list)) and len(dim) == 0):
+        dim = tuple(range(len(input_val.shape)))
+
     layer = ctx.net.add_reduce(
         input_val,
         trt.ReduceOperation.MAX,
@@ -51,7 +54,7 @@ def sum(
     ):
         input_val = cast_trt_tensor(ctx, input_val, trt.float32, name)
 
-    if dim is None:
+    if dim is None or (isinstance(dim, (tuple, list)) and len(dim) == 0):
         dim = tuple(range(len(input_val.shape)))
 
     layer = ctx.net.add_reduce(
@@ -169,7 +172,7 @@ def mean(
     ):
         input_val = cast_trt_tensor(ctx, input_val, trt.float32, name)
 
-    if dim is None:
+    if dim is None or (isinstance(dim, (tuple, list)) and len(dim) == 0):
         dim = tuple(range(len(input_val.shape)))
 
     layer = ctx.net.add_reduce(

tests/py/dynamo/conversion/test_amax_aten.py

@@ -30,6 +30,7 @@ def forward(self, x):
 
     @parameterized.expand(
         [
+            ((1, 2, 4), [], True),
             ((3, 2, 4), [1], True),
             ((2, 1, 4, 5), [0, 3], True),
             ((2, 3, 4, 5), [0, 1, 2, 3], False),
@@ -72,6 +73,7 @@ def forward(self, x):
 
     @parameterized.expand(
         [
+            ((1, 2, 4), [], True, torch.int, 0, 5),
             ((3, 2, 4), [1], True, torch.int, 0, 5),
             ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10),
             ((2, 3, 4, 5), [0, 1, 2, 3], False, torch.int32, -5, 0),

tests/py/dynamo/conversion/test_max_aten.py

@@ -9,9 +9,9 @@
 class TestMaxConverter(DispatchTestCase):
     @parameterized.expand(
         [
+            ((1, 2),),
             ((3, 2, 4),),
             ((2, 3, 4, 5),),
-            ((2, 3, 4, 5),),
             ((6, 7, 5, 4, 5),),
         ]
     )

tests/py/dynamo/conversion/test_min_aten.py

@@ -9,9 +9,9 @@
 class TestMinConverter(DispatchTestCase):
     @parameterized.expand(
         [
+            ((1, 2),),
             ((3, 2, 4),),
             ((2, 3, 4, 5),),
-            ((2, 3, 4, 5),),
             ((6, 7, 5, 4, 5),),
         ]
     )

tests/py/dynamo/conversion/test_prod_aten.py

@@ -9,9 +9,9 @@
 class TestProdConverter(DispatchTestCase):
     @parameterized.expand(
         [
+            ((1, 2),),
             ((3, 2, 4),),
             ((2, 3, 4, 5),),
-            ((2, 3, 4, 5),),
             ((6, 7, 5, 4, 5),),
         ]
     )

tests/py/dynamo/conversion/test_sum_aten.py

@@ -9,9 +9,9 @@
 class TestSumConverter(DispatchTestCase):
     @parameterized.expand(
         [
+            ((1, 2),),
             ((3, 2, 4),),
             ((2, 3, 4, 5),),
-            ((2, 3, 4, 5),),
             ((6, 7, 5, 4, 5),),
         ]
     )
@@ -51,6 +51,7 @@ def forward(self, x):
 
     @parameterized.expand(
         [
+            ((1, 2, 4), [], True),
             ((3, 2, 4), [1], True),
             ((2, 1, 4, 5), None, True),
             ((2, 3, 4, 5), [0, 1, 2, 3], False),
@@ -93,6 +94,7 @@ def forward(self, x):
 
     @parameterized.expand(
         [
+            ((1, 2, 4), [], True, torch.int, 0, 5),
             ((3, 2, 4), [1], True, torch.int, 0, 5),
             ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10),
             ((2, 3, 4, 5), None, False, torch.int32, -5, 0),