fix dim bugs

pytorch · Oct 5, 2023 · d4f1451 · d4f1451
1 parent 016f266
commit d4f1451
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 4 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/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),
     )
 

diff --git a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py b/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(

diff --git a/tests/py/dynamo/conversion/test_amax_aten.py b/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),

diff --git a/tests/py/dynamo/conversion/test_sum_aten.py b/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),