chunk_validator

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -692,7 +692,7 @@ def aten_ops_softmax(
 
 @dynamo_tensorrt_converter(
     torch.ops.aten.split.Tensor,
-    capability_validator=has_static_shapes_in_args([1]),
+    capability_validator= (has_static_shapes_in_args([0]) and has_static_shapes_in_args([1])),
     supports_dynamic_shapes=True,
 )
 @dynamo_tensorrt_converter(
@@ -903,7 +903,10 @@ def aten_ops_slice(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.chunk.default)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.chunk.default,
+    supports_dynamic_shapes=True,
+)
 @enforce_tensor_types(
     {
         0: (TRTTensor,),

tests/py/dynamo/conversion/test_chunk_aten.py

@@ -1,6 +1,7 @@
 import torch
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
 
 from .harness import DispatchTestCase
 
@@ -27,6 +28,7 @@ def forward(self, input):
         self.run_test(
             TestChunk(),
             input,
+            use_dynamo_tracer=True,
         )
 
     @parameterized.expand(
@@ -51,6 +53,7 @@ def forward(self, input):
         self.run_test(
             TestChunk(),
             input,
+            use_dynamo_tracer=True,
         )
 
     @parameterized.expand(
@@ -75,8 +78,102 @@ def forward(self, input):
         self.run_test(
             TestChunk(),
             input,
+            use_dynamo_tracer=True,
         )
 
+#######################Dynamic cases################
+
+    # @parameterized.expand(
+    #     [
+    #         ((1,), (1,), (3,), 3, 0),
+    #         ((3,), (3,), (4,), 3, 0),
+    #         ((4,), (4,), (6,), 3, 0),
+    #         ((6,), (6,), (9,), 3, 0),
+    #         ((3,), (3,), (4,), 1, -1),
+    #         ((3,), (3,), (4,), 3, -1),
+    #         ((3,), (3,), (4,), 4, -1),
+    #     ]
+    # )
+    # def test_chunk_1D(self, min_shape, opt_shape, max_shape, chunks, dim):
+    #     class TestChunk(torch.nn.Module):
+    #         def forward(self, input):
+    #             out = torch.ops.aten.chunk.default(input, chunks, dim)
+    #             return out
+
+    #     input_specs = [
+    #         Input(
+    #             min_shape=min_shape,
+    #             opt_shape=opt_shape,
+    #             max_shape=max_shape,
+    #         ),
+    #     ]
+    #     self.run_test_with_dynamic_shape(
+    #         TestChunk(),
+    #         input_specs,
+    #         use_dynamo_tracer = True,
+    #     )
+
+    # @parameterized.expand(
+    #     [
+    #         ((3, 4), (3, 4), (4, 4), 1, 0),
+    #         ((3, 4), (3, 4), (4, 4), 3, 0),
+    #         ((3, 4), (3, 4), (4, 4), 4, 0),
+    #         ((3, 4), (3, 4), (4, 4), 2, -2),
+    #         ((3, 4), (3, 4), (4, 4), 6, -2),
+    #         ((3, 4), (3, 4), (4, 4), 3, 1),
+    #         ((3, 4), (3, 4), (4, 4), 4, 1),
+    #         ((3, 4), (3, 4), (4, 4), 5, -1),
+    #     ]
+    # )
+    # def test_chunk_2D(self, min_shape, opt_shape, max_shape, chunks, dim):
+    #     class TestChunk(torch.nn.Module):
+    #         def forward(self, input):
+    #             out = torch.ops.aten.chunk.default(input, chunks, dim)
+    #             return out
+
+    #     input_specs = [
+    #         Input(
+    #             min_shape=min_shape,
+    #             opt_shape=opt_shape,
+    #             max_shape=max_shape,
+    #         ),
+    #     ]
+    #     self.run_test_with_dynamic_shape(
+    #         TestChunk(),
+    #         input_specs,
+    #         use_dynamo_tracer = True,
+    #     )
+
+    # @parameterized.expand(
+    #     [
+    #         ((3, 4, 2), (3, 4, 2), (4, 4, 2), 1, 0),
+    #         ((3, 4, 2), (3, 4, 2), (4, 4, 2), 3, -3),
+    #         ((3, 4, 2), (3, 4, 2), (4, 4, 2), 3, 1),
+    #         ((3, 4, 2), (3, 4, 2), (4, 4, 2), 4, 1),
+    #         ((3, 4, 2), (3, 4, 2), (4, 4, 2), 6, -2),
+    #         ((3, 4, 2), (3, 4, 2), (4, 4, 2), 1, 2),
+    #         ((3, 4, 2), (3, 4, 2), (4, 4, 2), 3, -1),
+    #         ((3, 4, 2), (3, 4, 2), (4, 4, 2), 4, -1),
+    #     ]
+    # )
+    # def test_chunk_3D(self, min_shape, opt_shape, max_shape, chunks, dim):
+    #     class TestChunk(torch.nn.Module):
+    #         def forward(self, input):
+    #             out = torch.ops.aten.chunk.default(input, chunks, dim)
+    #             return out
+
+    #     input_specs = [
+    #         Input(
+    #             min_shape=min_shape,
+    #             opt_shape=opt_shape,
+    #             max_shape=max_shape,
+    #         ),
+    #     ]
+    #     self.run_test_with_dynamic_shape(
+    #         TestChunk(),
+    #         input_specs,
+    #         use_dynamo_tracer = True,
+    #     )
 
 if __name__ == "__main__":
     run_tests()