Update on "[dynamo] support custom __getattr__ on torch.nn.Modules"

davidberard98 · davidberard98 · commit 1c78962d9138 · 2023-02-13T19:57:03.000-08:00
**Summary**: torch.nn.Module implementations previously did not support custom implementations of `__getattr__`; if a torch.nn.Module subclass implemented `__getattr__` and we tried to access an attribute that was expected to be present in `__getattr__`, dynamo would not check `__getattr__` and would error out with an AttributeError. This PR copies the functionality from UserDefinedObjectVariable into torch.nn.Module so that it also supports `__getattr__` Example of a module which previously would fail: ```python class MyMod(torch.nn.Module): def __init__(self): super().__init__() self.custom_dict = {"queue": [torch.rand((2, 2)) for _ in range(3)]} self.other_attr = torch.rand((2, 2)) def __getattr__(self, name): custom_dict = self.custom_dict if name in custom_dict: return custom_dict[name] return super().__getattr__(name) def forward(self, x): return x @ self.other_attr + self.queue[-1] ``` cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 @EikanWang jgong5 @Guobing-Chen @XiaobingSuper zhuhaozhe blzheng @Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py
@@ -1033,7 +1033,6 @@ def fn(obj, x):
         cnts = torch._dynamo.testing.CompileCounter()
         opt_fn = torch._dynamo.optimize(cnts)(fn)
         self.assertTrue(same(opt_fn(obj, x), fn(obj, x)))
-        self.assertTrue(cnts.frame_count <= 2)
 
     def test_nn_module_getattr(self):
         class MyMod(torch.nn.Module):
@@ -1081,7 +1080,6 @@ def fn(mod, x):
         cnts = torch._dynamo.testing.CompileCounter()
         opt_fn = torch._dynamo.optimize(cnts)(fn)
         self.assertTrue(same(opt_fn(mod, x), fn(mod, x)))
-        self.assertTrue(cnts.frame_count <= 2)
 
     def test_user_property(self):
         class MyConfig:
