[dynamo] Inline the getattr of fx graph and proxy graph (#128172)

anijain2305 · pytorchmergebot · commit 662a78f957fb · 2024-06-07T17:14:58.000Z
Pull Request resolved: #128172 Approved by: https://github.com/yanboliang ghstack dependencies: #128001, #126578, #128158
diff --git a/test/dynamo/test_inline_inbuilt_nn_modules.py b/test/dynamo/test_inline_inbuilt_nn_modules.py
@@ -6,6 +6,7 @@
 try:
     from . import (
         test_aot_autograd,
+        test_export,
         test_functions,
         test_higher_order_ops,
         test_misc,
@@ -14,6 +15,7 @@
     )
 except ImportError:
     import test_aot_autograd
+    import test_export
     import test_functions
     import test_higher_order_ops
     import test_misc
@@ -50,6 +52,7 @@ def make_inline_inbuilt_nn_modules_cls(cls):
     test_higher_order_ops.HigherOrderOpTests,
     test_higher_order_ops.FuncTorchHigherOrderOpTests,
     test_aot_autograd.AotAutogradFallbackTests,
+    test_export.ExportTests,
     # test_repros.ReproTests,
 ]
 for test in tests:
diff --git a/torch/_dynamo/trace_rules.py b/torch/_dynamo/trace_rules.py
@@ -3231,6 +3231,8 @@ def _module_dir(m: types.ModuleType):
     "torch.cuda.amp.autocast_mode",
     "torch.distributions",
     "torch.fx._pytree",
+    "torch.fx._symbolic_trace",
+    "torch.fx.experimental.proxy_tensor",
     "torch.fx.passes.shape_prop",
     "torch.nn",
     "torch.random",
