pytorch

1
# Owner(s): ["module: fx"]
2

3
from __future__ import annotations    # type: ignore[attr-defined]
4
import torch
5
import typing
6
from torch.fx import symbolic_trace
7

8
class A:
9
    def __call__(self, x: torch.Tensor):
10
        return torch.add(x, x)
11

12
# No forward references
13
class M1(torch.nn.Module):
14
    def forward(self, x: torch.Tensor, a: A) -> torch.Tensor:
15
        return a(x)
16

17
# Forward references
18
class M2(torch.nn.Module):
19
    def forward(self, x: torch.Tensor, a: A) -> torch.Tensor:
20
        return a(x)
21

22
# Non-torch annotation with no internal forward references
23
class M3(torch.nn.Module):
24
    def forward(self, x: typing.List[torch.Tensor], a: A) -> torch.Tensor:
25
        return a(x[0])
26

27
# Non-torch annotation with internal forward references
28
class M4(torch.nn.Module):
29
    def forward(self, x: typing.List[torch.Tensor], a: A) -> torch.Tensor:
30
        return a(x[0])
31

32
x = torch.rand(2, 3)
33

34
ref = torch.add(x, x)
35

36
traced1 = symbolic_trace(M1())
37
res1 = traced1(x, A())
38
assert torch.all(torch.eq(ref, res1))
39

40
traced2 = symbolic_trace(M2())
41
res2 = traced2(x, A())
42
assert torch.all(torch.eq(ref, res2))
43

44
traced3 = symbolic_trace(M3())
45
res3 = traced3([x], A())
46
assert torch.all(torch.eq(ref, res3))
47

48
traced4 = symbolic_trace(M4())
49
res4 = traced4([x], A())
50
assert torch.all(torch.eq(ref, res4))
51