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