torchdynamo support self.modules() for nn_module (pytorch#88695)

This PR allows models to call self.modules() during dynamo tracing.

Pull Request resolved: pytorch#88695
Approved by: https://github.com/voznesenskym

test/dynamo/test_repros.py

@@ -1792,6 +1792,26 @@ def fn(x):
         res = opt_fn(a)
         self.assertTrue(same(ref, res))
 
+    def test_modules(self):
+        class Foo(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.fc = torch.nn.Linear(4, 3)
+
+            def forward(self, inp):
+                res = torch.zeros(3, 3)
+                for mod in self.modules():
+                    res += self.fc(inp)
+                return res
+
+        mod = Foo()
+        args = (torch.ones(3, 4),)
+        cnt = torch._dynamo.testing.CompileCounter()
+        opt_mod = torch._dynamo.optimize(cnt, nopython=True)(mod)
+        self.assertTrue(same(mod(*args), opt_mod(*args)))
+        self.assertEqual(cnt.op_count, 5)
+        self.assertEqual(cnt.frame_count, 1)
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests

torch/_dynamo/guards.py

@@ -543,7 +543,7 @@ def __init__(self, expr_to_tensor_ref, id_to_name_map):
         self.id_to_name_map = id_to_name_map
 
     def _print_Symbol(self, expr) -> str:
-        assert isinstance(expr, sympy.core.symbol.Symbol)
+        assert isinstance(expr, sympy.Symbol)
         if expr == 0:
             return "0"
         if expr == 1:

torch/_dynamo/variables/nn_module.py

@@ -337,6 +337,8 @@ def named_embed(name, obj):
             ):
                 result.append(named_embed(name, submod))
             return ListIteratorVariable(result, mutable_local=MutableLocal(), **options)
+        elif name == "modules":
+            return wrap_values(module.named_modules())
         elif name == "parameters":
             return wrap_values(module.named_parameters(**get_kwargs("recurse")))
         elif name == "values":