clean flatten converter and add tests

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -483,7 +483,7 @@ def aten_ops_slice(
     {
         0: (TRTTensor,),
     }
-)
+)  # type: ignore[misc]
 def aten_ops_permute(
     ctx: ConversionContext,
     target: Target,
@@ -1371,7 +1371,7 @@ def conv_param_validator(conv_node: Node) -> bool:
         1: (np.ndarray, torch.Tensor, TRTTensor),
         2: (np.ndarray, torch.Tensor, TRTTensor),
     }
-)
+)  # type: ignore[misc]
 def aten_ops_convolution(
     ctx: ConversionContext,
     target: Target,
@@ -1551,23 +1551,3 @@ def aten_ops_reshape(
         input=args[0],
         shape=args[1],
     )
-
-
-# # TODO: need tests for this converter
-# @dynamo_tensorrt_converter(torch.ops.aten.flatten.using_ints)  # type: ignore[misc]
-# def aten_ops_flatten(
-#     ctx: ConversionContext,
-#     target: Target,
-#     args: Tuple[Argument, ...],
-#     kwargs: Dict[str, Argument],
-#     name: str,
-# ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-#     return impl.shuffle.flatten(
-#         ctx,
-#         target,
-#         SourceIR.ATEN,
-#         name,
-#         input=args[0],
-#         start_dim=args_bounds_check(args, 1, 0),
-#         end_dim=args_bounds_check(args, 2, -1),
-#     )

tests/py/dynamo/conversion/test_converter_utils.py

@@ -1,7 +1,11 @@
 import numpy as np
 import torch
+from parameterized import parameterized
 from torch.testing._internal.common_utils import TestCase, run_tests
-from torch_tensorrt.dynamo.conversion.converter_utils import enforce_tensor_types
+from torch_tensorrt.dynamo.conversion.converter_utils import (
+    enforce_tensor_types,
+    flatten_dims,
+)
 from torch_tensorrt.fx.types import TRTTensor
 
 from ..testing_utilities import DECIMALS_OF_AGREEMENT, lower_graph_testing
@@ -37,5 +41,39 @@ def test_invalid_invocation_type(self):
             enforce_tensor_types({0: (int, bool)})
 
 
+class TestFlattenDimsEnforcement(TestCase):
+    @parameterized.expand(
+        [
+            ((1, 2), 0, 0, (1, 2)),
+            ((1, 2), 0, 1, (2,)),
+            ((2, 3, 4), 1, 2, (2, 12)),
+            ((2, 3, 4), 0, 1, (6, 4)),
+            ((2, 3, 4), -3, 2, (24,)),
+            ((2, 3, 4, 5), 0, -2, (24, 5)),
+            ((2, 3, 4, 5), -4, -1, (120,)),
+        ]
+    )
+    def test_numpy_array(self, input_shape, start_dim, end_dim, true_shape):
+        inputs = np.random.randn(*input_shape)
+        new_shape = flatten_dims(inputs, start_dim, end_dim)
+        self.assertEqual(new_shape, true_shape)
+
+    @parameterized.expand(
+        [
+            ((1, 2), 0, 0, (1, 2)),
+            ((1, 2), 0, 1, (2,)),
+            ((2, 3, 4), 1, 2, (2, 12)),
+            ((2, 3, 4), 0, 1, (6, 4)),
+            ((2, 3, 4), -3, 2, (24,)),
+            ((2, 3, 4, 5), 0, -2, (24, 5)),
+            ((2, 3, 4, 5), -4, -1, (120,)),
+        ]
+    )
+    def test_torch_tensor(self, input_shape, start_dim, end_dim, true_shape):
+        inputs = torch.randn(input_shape)
+        new_shape = flatten_dims(inputs, start_dim, end_dim)
+        self.assertEqual(new_shape, true_shape)
+
+
 if __name__ == "__main__":
     run_tests()