TensorRT/tests/py/dynamo/conversion/test_embedding_aten.py at 87c64d3b9145f176be26dbe8c133dc2972caa576 · 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
110
111
112
113
114
115
116
117
118
119
import torch
from parameterized import param, parameterized
from torch.testing._internal.common_utils import run_tests
from torch_tensorrt import Input

from .harness import DispatchTestCase


class TestEmbeddingConverter(DispatchTestCase):
    @parameterized.expand(
        [
            param(
                test_name="1d_indices",
                indices_tensor=torch.tensor([3, 1, 2], dtype=torch.int32),
                weights_tensor=torch.randn((5, 10), dtype=torch.float32),
                sparse=False,
            ),
            param(
                test_name="2d_indices",
                indices_tensor=torch.tensor([[3, 1, 2], [4, 1, 3]], dtype=torch.int32),
                weights_tensor=torch.randn((5, 10), dtype=torch.float32),
                sparse=True,
            ),
            param(
                test_name="3d_indices",
                indices_tensor=torch.tensor(
                    [[[0, 1], [2, 3]], [[3, 4], [4, 0]]], dtype=torch.int32
                ),
                weights_tensor=torch.randn((5, 10), dtype=torch.float32),
                sparse=True,
            ),
            # int64 indices - TensorRT now supports int64 for gather operations
            param(
                test_name="1d_indices_int64",
                indices_tensor=torch.tensor([3, 1, 2], dtype=torch.int64),
                weights_tensor=torch.randn((5, 10), dtype=torch.float32),
                sparse=False,
            ),
            param(
                test_name="2d_indices_int64",
                indices_tensor=torch.tensor([[3, 1, 2], [4, 1, 3]], dtype=torch.int64),
                weights_tensor=torch.randn((5, 10), dtype=torch.float32),
                sparse=True,
            ),
            param(
                test_name="3d_indices_int64",
                indices_tensor=torch.tensor(
                    [[[0, 1], [2, 3]], [[3, 4], [4, 0]]], dtype=torch.int64
                ),
                weights_tensor=torch.randn((5, 10), dtype=torch.float32),
                sparse=True,
            ),
        ]
    )
    def test_embedding(
        self,
        test_name,
        indices_tensor,
        weights_tensor,
        padding_idx=-1,
        max_norm=None,
        norm_type=2.0,
        scale_grad_by_freq=None,
        sparse=False,
    ):
        class TestEmbedding(torch.nn.Module):
            def forward(self, indices, weights):
                return torch.ops.aten.embedding.default(
                    weights,
                    indices,
                    padding_idx,
                    scale_grad_by_freq,
                    sparse,
                )

        self.run_test(
            TestEmbedding(),
            inputs=[indices_tensor, weights_tensor],
        )

    def test_embedding_with_dynamic_shape_four_dimensions(
        self,
        padding_idx=-1,
        max_norm=None,
        norm_type=2.0,
        scale_grad_by_freq=None,
        sparse=None,
    ):
        class TestEmbedding(torch.nn.Module):
            def forward(self, input, weights):
                return torch.ops.aten.embedding.default(
                    weights,
                    input,
                    padding_idx,
                    scale_grad_by_freq,
                    sparse,
                )

        input_specs = [
            Input(
                shape=(-1, -1, -1, -1),
                dtype=torch.int32,
                shape_ranges=[((1, 1, 1, 1), (2, 3, 4, 5), (2, 3, 10, 10))],
            ),
            Input(
                shape=(-1, -1),
                dtype=torch.float32,
                shape_ranges=[((1, 1), (2, 3), (2, 3))],
            ),
        ]

        self.run_test_with_dynamic_shape(
            TestEmbedding(),
            input_specs,
        )


if __name__ == "__main__":
    run_tests()