Skip to content

Commit 1f10647

Browse files
chohk88chohk88
authored
feat: dynamic shape support for tan, sinh, cosh, asin and acos (#2941)
1 parent 7b825f5 commit 1f10647

File tree

6 files changed

+257
-5
lines changed

6 files changed

+257
-5
lines changed

‎py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1511,7 +1511,7 @@ def aten_ops_cos(
15111511
)
15121512

15131513

1514-
@dynamo_tensorrt_converter(torch.ops.aten.tan.default)
1514+
@dynamo_tensorrt_converter(torch.ops.aten.tan.default, supports_dynamic_shapes=True)
15151515
def aten_ops_tan(
15161516
ctx: ConversionContext,
15171517
target: Target,
@@ -1528,7 +1528,7 @@ def aten_ops_tan(
15281528
)
15291529

15301530

1531-
@dynamo_tensorrt_converter(torch.ops.aten.sinh.default)
1531+
@dynamo_tensorrt_converter(torch.ops.aten.sinh.default, supports_dynamic_shapes=True)
15321532
def aten_ops_sinh(
15331533
ctx: ConversionContext,
15341534
target: Target,
@@ -1545,7 +1545,7 @@ def aten_ops_sinh(
15451545
)
15461546

15471547

1548-
@dynamo_tensorrt_converter(torch.ops.aten.cosh.default)
1548+
@dynamo_tensorrt_converter(torch.ops.aten.cosh.default, supports_dynamic_shapes=True)
15491549
def aten_ops_cosh(
15501550
ctx: ConversionContext,
15511551
target: Target,
@@ -1562,7 +1562,7 @@ def aten_ops_cosh(
15621562
)
15631563

15641564

1565-
@dynamo_tensorrt_converter(torch.ops.aten.asin.default)
1565+
@dynamo_tensorrt_converter(torch.ops.aten.asin.default, supports_dynamic_shapes=True)
15661566
def aten_ops_asin(
15671567
ctx: ConversionContext,
15681568
target: Target,
@@ -1579,7 +1579,7 @@ def aten_ops_asin(
15791579
)
15801580

15811581

1582-
@dynamo_tensorrt_converter(torch.ops.aten.acos.default)
1582+
@dynamo_tensorrt_converter(torch.ops.aten.acos.default, supports_dynamic_shapes=True)
15831583
def aten_ops_acos(
15841584
ctx: ConversionContext,
15851585
target: Target,

‎tests/py/dynamo/conversion/test_acos_aten.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import torch.nn as nn
33
from parameterized import parameterized
44
from torch.testing._internal.common_utils import run_tests
5+
from torch_tensorrt import Input
56

67
from .harness import DispatchTestCase
78

@@ -44,6 +45,56 @@ def forward(self, input):
4445
inputs,
4546
)
4647

48+
@parameterized.expand(
49+
[
50+
(
51+
"3d_dim_dtype_int32",
52+
(3, 2, 1),
53+
(3, 2, 3),
54+
(3, 3, 4),
55+
torch.int32,
56+
torch.float32,
57+
),
58+
(
59+
"2d_dim_dtype_float16",
60+
(1, 1),
61+
(2, 2),
62+
(4, 4),
63+
torch.float16,
64+
torch.float16,
65+
),
66+
(
67+
"3d_dim_dtype_float",
68+
(1, 1, 1),
69+
(1, 2, 3),
70+
(3, 3, 3),
71+
torch.float,
72+
torch.float,
73+
),
74+
]
75+
)
76+
def test_acos_dynamic_shape(
77+
self, _, min_shape, opt_shape, max_shape, type, output_type
78+
):
79+
class acos(nn.Module):
80+
def forward(self, input):
81+
return torch.ops.aten.acos.default(input)
82+
83+
input_specs = [
84+
Input(
85+
min_shape=min_shape,
86+
opt_shape=opt_shape,
87+
max_shape=max_shape,
88+
dtype=type,
89+
),
90+
]
91+
92+
self.run_test_with_dynamic_shape(
93+
acos(),
94+
input_specs,
95+
output_dtypes=[output_type],
96+
)
97+
4798

4899
if __name__ == "__main__":
49100
run_tests()

‎tests/py/dynamo/conversion/test_asin_aten.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import torch.nn as nn
33
from parameterized import parameterized
44
from torch.testing._internal.common_utils import run_tests
5+
from torch_tensorrt import Input
56

67
from .harness import DispatchTestCase
78

@@ -44,6 +45,55 @@ def forward(self, input):
4445
inputs,
4546
)
4647

48+
@parameterized.expand(
49+
[
50+
(
51+
"3d_dim_dtype_int32",
52+
(3, 2, 1),
53+
(3, 2, 3),
54+
(3, 3, 4),
55+
torch.int32,
56+
torch.float32,
57+
),
58+
(
59+
"2d_dim_dtype_float16",
60+
(1, 1),
61+
(2, 2),
62+
(4, 4),
63+
torch.float16,
64+
torch.float16,
65+
),
66+
(
67+
"3d_dim_dtype_float",
68+
(1, 1, 1),
69+
(1, 2, 3),
70+
(3, 3, 3),
71+
torch.float,
72+
torch.float,
73+
),
74+
]
75+
)
76+
def test_asin_dynamic_shape(
77+
self, _, min_shape, opt_shape, max_shape, type, output_type
78+
):
79+
class asin(nn.Module):
80+
def forward(self, input):
81+
return torch.ops.aten.asin.default(input)
82+
83+
input_specs = [
84+
Input(
85+
min_shape=min_shape,
86+
opt_shape=opt_shape,
87+
max_shape=max_shape,
88+
dtype=type,
89+
),
90+
]
91+
self.run_test_with_dynamic_shape(
92+
asin(),
93+
input_specs,
94+
output_dtypes=[output_type],
95+
)
96+
4797

4898
if __name__ == "__main__":
4999
run_tests()

‎tests/py/dynamo/conversion/test_cosh_aten.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import torch.nn as nn
33
from parameterized import parameterized
44
from torch.testing._internal.common_utils import run_tests
5+
from torch_tensorrt import Input
56

67
from .harness import DispatchTestCase
78

@@ -44,6 +45,55 @@ def forward(self, input):
4445
inputs,
4546
)
4647

48+
@parameterized.expand(
49+
[
50+
(
51+
"3d_dim_dtype_int32",
52+
(3, 2, 1),
53+
(3, 2, 3),
54+
(3, 3, 4),
55+
torch.int32,
56+
torch.float32,
57+
),
58+
(
59+
"2d_dim_dtype_float16",
60+
(1, 1),
61+
(2, 2),
62+
(4, 4),
63+
torch.float16,
64+
torch.float16,
65+
),
66+
(
67+
"3d_dim_dtype_float",
68+
(1, 1, 1),
69+
(1, 2, 3),
70+
(3, 3, 3),
71+
torch.float,
72+
torch.float,
73+
),
74+
]
75+
)
76+
def test_cosh_dynamic_shape(
77+
self, _, min_shape, opt_shape, max_shape, type, output_type
78+
):
79+
class cosh(nn.Module):
80+
def forward(self, input):
81+
return torch.ops.aten.cosh.default(input)
82+
83+
input_specs = [
84+
Input(
85+
min_shape=min_shape,
86+
opt_shape=opt_shape,
87+
max_shape=max_shape,
88+
dtype=type,
89+
),
90+
]
91+
self.run_test_with_dynamic_shape(
92+
cosh(),
93+
input_specs,
94+
output_dtypes=[output_type],
95+
)
96+
4797

4898
if __name__ == "__main__":
4999
run_tests()

‎tests/py/dynamo/conversion/test_sinh_aten.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import torch.nn as nn
33
from parameterized import parameterized
44
from torch.testing._internal.common_utils import run_tests
5+
from torch_tensorrt import Input
56

67
from .harness import DispatchTestCase
78

@@ -44,6 +45,55 @@ def forward(self, input):
4445
inputs,
4546
)
4647

48+
@parameterized.expand(
49+
[
50+
(
51+
"3d_dim_dtype_int32",
52+
(3, 2, 1),
53+
(3, 2, 3),
54+
(3, 3, 4),
55+
torch.int32,
56+
torch.float32,
57+
),
58+
(
59+
"2d_dim_dtype_float16",
60+
(1, 1),
61+
(2, 2),
62+
(4, 4),
63+
torch.float16,
64+
torch.float16,
65+
),
66+
(
67+
"3d_dim_dtype_float",
68+
(1, 1, 1),
69+
(1, 2, 3),
70+
(3, 3, 3),
71+
torch.float,
72+
torch.float,
73+
),
74+
]
75+
)
76+
def test_sinh_dynamic_shape(
77+
self, _, min_shape, opt_shape, max_shape, type, output_type
78+
):
79+
class sinh(nn.Module):
80+
def forward(self, input):
81+
return torch.ops.aten.sinh.default(input)
82+
83+
input_specs = [
84+
Input(
85+
min_shape=min_shape,
86+
opt_shape=opt_shape,
87+
max_shape=max_shape,
88+
dtype=type,
89+
),
90+
]
91+
self.run_test_with_dynamic_shape(
92+
sinh(),
93+
input_specs,
94+
output_dtypes=[output_type],
95+
)
96+
4797

4898
if __name__ == "__main__":
4999
run_tests()

‎tests/py/dynamo/conversion/test_tan_aten.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import torch.nn as nn
33
from parameterized import parameterized
44
from torch.testing._internal.common_utils import run_tests
5+
from torch_tensorrt import Input
56

67
from .harness import DispatchTestCase
78

@@ -44,6 +45,56 @@ def forward(self, input):
4445
inputs,
4546
)
4647

48+
@parameterized.expand(
49+
[
50+
(
51+
"3d_dim_dtype_int32",
52+
(3, 2, 1),
53+
(3, 2, 3),
54+
(3, 3, 4),
55+
torch.int32,
56+
torch.float32,
57+
),
58+
(
59+
"2d_dim_dtype_float16",
60+
(1, 1),
61+
(2, 2),
62+
(4, 4),
63+
torch.float16,
64+
torch.float16,
65+
),
66+
(
67+
"3d_dim_dtype_float",
68+
(1, 1, 1),
69+
(1, 2, 3),
70+
(3, 3, 3),
71+
torch.float,
72+
torch.float,
73+
),
74+
]
75+
)
76+
def test_tan_dynamic_shape(
77+
self, _, min_shape, opt_shape, max_shape, type, output_type
78+
):
79+
class tan(nn.Module):
80+
def forward(self, input):
81+
return torch.ops.aten.tan.default(input)
82+
83+
input_specs = [
84+
Input(
85+
min_shape=min_shape,
86+
opt_shape=opt_shape,
87+
max_shape=max_shape,
88+
dtype=type,
89+
),
90+
]
91+
92+
self.run_test_with_dynamic_shape(
93+
tan(),
94+
input_specs,
95+
output_dtypes=[output_type],
96+
)
97+
4798

4899
if __name__ == "__main__":
49100
run_tests()

0 commit comments

Comments
 (0)