feat: dynamic shapes support for sqrt and copy (#2889)

Author: chohk88

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1289,7 +1289,7 @@ def aten_ops_log1p(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.sqrt.default)
+@dynamo_tensorrt_converter(torch.ops.aten.sqrt.default, supports_dynamic_shapes=True)
 def aten_ops_sqrt(
     ctx: ConversionContext,
     target: Target,
@@ -2938,7 +2938,7 @@ def aten_ops_trunc(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.copy.default)
+@dynamo_tensorrt_converter(torch.ops.aten.copy.default, supports_dynamic_shapes=True)
 @enforce_tensor_types(
     {
         1: (TRTTensor,),

tests/py/dynamo/conversion/test_copy_aten.py

@@ -2,6 +2,7 @@
 import torch.nn as nn
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
 
 from .harness import DispatchTestCase
 
@@ -26,6 +27,67 @@ def forward(self, input, src):
             inputs,
         )
 
+    @parameterized.expand(
+        [
+            (
+                "1d_float32",
+                torch.float32,
+                (1,),
+                (2,),
+                (3,),
+                False,
+            ),
+            (
+                "2d_float32",
+                torch.float32,
+                (1, 1),
+                (2, 2),
+                (3, 3),
+                False,
+            ),
+            (
+                "3d_float32",
+                torch.float32,
+                (1, 1, 1),
+                (2, 2, 2),
+                (3, 3, 3),
+                True,
+            ),
+            (
+                "4d_float32",
+                torch.float32,
+                (1, 1, 1, 1),
+                (2, 2, 2, 2),
+                (3, 3, 3, 3),
+                True,
+            ),
+        ]
+    )
+    def test_dynamic_shape_copy_float(self, *args):
+        class Copy(nn.Module):
+            def forward(self, input, src):
+                return torch.ops.aten.copy.default(input, src, args[5])
+
+        input_specs = [
+            Input(
+                min_shape=args[2],
+                opt_shape=args[3],
+                max_shape=args[4],
+                dtype=args[1],
+            ),
+            Input(
+                min_shape=args[2],
+                opt_shape=args[3],
+                max_shape=args[4],
+                dtype=args[1],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            Copy(),
+            input_specs,
+        )
+
 
 if __name__ == "__main__":
     run_tests()

tests/py/dynamo/conversion/test_sqrt_aten.py

@@ -2,6 +2,7 @@
 import torch.nn as nn
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
 
 from .harness import DispatchTestCase
 
@@ -44,6 +45,100 @@ def forward(self, input):
             inputs,
         )
 
+    @parameterized.expand(
+        [
+            (
+                "1d_float32",
+                torch.float32,
+                (1,),
+                (2,),
+                (3,),
+            ),
+            (
+                "2d_float32",
+                torch.float32,
+                (1, 1),
+                (2, 2),
+                (3, 3),
+            ),
+            (
+                "3d_float32",
+                torch.float32,
+                (1, 1, 1),
+                (2, 2, 2),
+                (3, 3, 3),
+            ),
+            (
+                "4d_float32",
+                torch.float32,
+                (1, 1, 1, 1),
+                (2, 2, 2, 2),
+                (3, 3, 3, 3),
+            ),
+        ]
+    )
+    def test_dynamic_shape_sqrt_float(self, *args):
+        class sqrt(nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.sqrt.default(input)
+
+        input_specs = [
+            Input(
+                min_shape=args[2],
+                opt_shape=args[3],
+                max_shape=args[4],
+                dtype=args[1],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(sqrt(), input_specs)
+
+    @parameterized.expand(
+        [
+            (
+                "1d_int32",
+                torch.int,
+                (1,),
+                (2,),
+                (3,),
+            ),
+            (
+                "2d_int32",
+                torch.int32,
+                (1, 1),
+                (2, 2),
+                (3, 3),
+            ),
+            (
+                "3d_int32",
+                torch.int,
+                (1, 1, 1),
+                (2, 2, 2),
+                (3, 3, 3),
+            ),
+            (
+                "4d_int32",
+                torch.int32,
+                (1, 1, 1, 1),
+                (2, 2, 2, 2),
+                (3, 3, 3, 3),
+            ),
+        ]
+    )
+    def test_dynamic_shape_sqrt_int(self, *args):
+        class sqrt(nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.sqrt.default(input)
+
+        input_specs = [
+            Input(
+                min_shape=args[2],
+                opt_shape=args[3],
+                max_shape=args[4],
+                dtype=args[1],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(sqrt(), input_specs)
+
 
 if __name__ == "__main__":
     run_tests()