xref: /aosp_15_r20/external/pytorch/test/jit/test_attr.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# Owner(s): ["oncall: jit"]
2
3from typing import NamedTuple, Tuple
4
5import torch
6from torch.testing import FileCheck
7from torch.testing._internal.jit_utils import JitTestCase
8
9
10if __name__ == "__main__":
11    raise RuntimeError(
12        "This test file is not meant to be run directly, use:\n\n"
13        "\tpython test/test_jit.py TESTNAME\n\n"
14        "instead."
15    )
16
17
18class TestGetDefaultAttr(JitTestCase):
19    def test_getattr_with_default(self):
20        class A(torch.nn.Module):
21            def __init__(self) -> None:
22                super().__init__()
23                self.init_attr_val = 1.0
24
25            def forward(self, x):
26                y = getattr(self, "init_attr_val")  # noqa: B009
27                w: list[float] = [1.0]
28                z = getattr(self, "missing", w)  # noqa: B009
29                z.append(y)
30                return z
31
32        result = A().forward(0.0)
33        self.assertEqual(2, len(result))
34        graph = torch.jit.script(A()).graph
35
36        # The "init_attr_val" attribute exists
37        FileCheck().check('prim::GetAttr[name="init_attr_val"]').run(graph)
38        # The "missing" attribute does not exist, so there should be no corresponding GetAttr in AST
39        FileCheck().check_not("missing").run(graph)
40        # instead the getattr call will emit the default value, which is a list with one float element
41        FileCheck().check("float[] = prim::ListConstruct").run(graph)
42
43    def test_getattr_named_tuple(self):
44        global MyTuple
45
46        class MyTuple(NamedTuple):
47            x: str
48            y: torch.Tensor
49
50        def fn(x: MyTuple) -> Tuple[str, torch.Tensor, int]:
51            return (
52                getattr(x, "x", "fdsa"),
53                getattr(x, "y", torch.ones((3, 3))),
54                getattr(x, "z", 7),
55            )
56
57        inp = MyTuple(x="test", y=torch.ones(3, 3) * 2)
58        ref = fn(inp)
59        fn_s = torch.jit.script(fn)
60        res = fn_s(inp)
61        self.assertEqual(res, ref)
62
63    def test_getattr_tuple(self):
64        def fn(x: Tuple[str, int]) -> int:
65            return getattr(x, "x", 2)
66
67        with self.assertRaisesRegex(RuntimeError, "but got a normal Tuple"):
68            torch.jit.script(fn)
69