TensorRT/tests/py/dynamo/conversion/test_cumsum_aten.py at 406a9e64cca85945ddf4db4e9ed61652b1ec1147 · pytorch/TensorRT · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import unittest

import torch
import torch.nn as nn
import torch_tensorrt
from parameterized import parameterized
from torch.testing._internal.common_utils import run_tests

from .harness import DispatchTestCase


@unittest.skipIf(
    torch_tensorrt.ENABLED_FEATURES.tensorrt_rtx,
    "cumsum is not supported on TensorRT-RTX (build_serialized_network returns None on Linux as well as Windows)",
)
class TestCumsumConverter(DispatchTestCase):
    @parameterized.expand(
        [
            ((1,), 0),
            ((2,), 0),
            ((3,), -1),
        ]
    )
    def test_cumsum_1D(self, shape, dim):
        class Cumsum(nn.Module):
            def forward(self, x):
                return torch.ops.aten.cumsum.default(x, dim)

        inputs = [torch.randn(shape)]
        self.run_test(
            Cumsum(),
            inputs,
            immutable_weights=False,
        )

    @parameterized.expand(
        [
            ((3, 1), 0),
            ((3, 1), 1),
            ((2, 3), -1),
            ((2, 3), -2),
        ]
    )
    def test_cumsum_2D(self, shape, dims):
        class Cumsum(nn.Module):
            def forward(self, x):
                return torch.ops.aten.cumsum.default(x, dims)

        inputs = [torch.randn(shape)]
        self.run_test(
            Cumsum(),
            inputs,
            immutable_weights=True,
        )

    @parameterized.expand(
        [
            ((2, 3, 3), 0),
            ((4, 2, 3), 1),
            ((1, 2, 3), 2),
            ((1, 2, 3), -1),
            ((1, 2, 3), -2),
        ]
    )
    def test_cumsum_3D(self, shape, dims):
        class Cumsum(nn.Module):
            def forward(self, x):
                return torch.ops.aten.cumsum.default(x, dims)

        inputs = [torch.randn(shape)]
        self.run_test(
            Cumsum(),
            inputs,
            immutable_weights=True,
        )

    @parameterized.expand(
        [
            ((1,), (2,), (3,), 0),
            ((1,), (2,), (3,), -1),
            ((2, 3), (2, 4), (2, 5), 0),
            ((2, 3), (3, 4), (4, 5), -1),
            ((1, 2, 2), (2, 2, 3), (3, 3, 3), 0),
            ((1, 2, 2), (2, 2, 3), (3, 2, 3), -2),
            ((1, 2, 2, 3), (2, 2, 3, 4), (3, 3, 4, 5), -3),
            ((1, 2, 2, 3), (2, 2, 3, 4), (3, 3, 4, 5), -2),
        ]
    )
    def test_cumsum_dynamic_shape(self, min_shape, opt_shape, max_shape, dims):
        class Cumsum(nn.Module):
            def forward(self, x):
                return torch.ops.aten.cumsum.default(x, dims)

        inputs = [
            torch_tensorrt.Input(
                min_shape=min_shape,
                opt_shape=opt_shape,
                max_shape=max_shape,
            ),
        ]
        self.run_test_with_dynamic_shape(
            Cumsum(),
            inputs,
            immutable_weights=False,
        )


if __name__ == "__main__":
    run_tests()