[export] Return more information from tracing context in graph capture. (pytorch#166775)

Summary:
as title, we should return an entire tracing_context object instead of fake_mode only, since tracing context should contain full set of information.

Test Plan:
pytest test/export/test_experimental.py

Pull Request resolved: pytorch#166775
Approved by: https://github.com/tugsbayasgalan

Author: zhxchen17

test/export/test_experimental.py

@@ -522,6 +522,25 @@ def forward(self, args_0):
         )
         self.assertEqual(ep(*inps), MyModel()(*inps))
 
+    def test_dynamo_graph_capture_full_tracing_context(self) -> None:
+        class Foo(torch.nn.Module):
+            def forward(self, x):
+                return x + x.shape[0]
+
+        foo = Foo()
+
+        def make_inputs(b: int):
+            ret = (torch.randn(b, 3),)
+            torch._dynamo.mark_dynamic(ret[0], 0)
+            return ret
+
+        trace_inputs = make_inputs(2)
+        gm = dynamo_graph_capture_for_export(foo)(*trace_inputs)
+        test_inputs = make_inputs(3)
+        self.assertEqual(gm(*test_inputs), foo(*test_inputs))
+        self.assertIsNotNone(gm.meta["tracing_context"].fake_mode)
+        self.assertEqual(len(gm.meta["tracing_context"].tensor_to_context), 1)
+
     @unittest.skipIf(not TEST_CUDA, "CUDA not available")
     def test_dynamo_graph_capture_fx_graph_annotate_overlap_pass(self):
         class DummyOp(torch.autograd.Function):

torch/_dynamo/functional_export.py

@@ -16,6 +16,7 @@
 from torch._dynamo.exc import UserErrorType
 from torch._dynamo.utils import dynamo_timed, get_metrics_context
 from torch._export.utils import _compiling_state_context
+from torch._guards import TracingContext
 from torch.export.dynamic_shapes import _RelaxedConstraint, Constraint
 from torch.fx import Node
 from torch.fx.experimental.proxy_tensor import make_fx
@@ -650,6 +651,10 @@ def inner(*args: Any, **kwargs: Any) -> Any:
         )
         assert out.backend_input is not None
         graph_module.meta["fake_mode"] = out.backend_input.fake_mode  # type: ignore[attr-defined]
+        graph_module.meta["fake_mode"].allow_non_fake_inputs = True
+        tracing_context = TracingContext(graph_module.meta["fake_mode"])
+        tracing_context.tensor_to_context = out.backend_input.tensor_to_context  # type: ignore[attr-defined]
+        graph_module.meta["tracing_context"] = tracing_context
         return graph_module
 
     return inner