change torch._dynamo.export(aten_graph=...) to allow pre_autograd tracing

test/dynamo/test_export.py

@@ -1717,6 +1717,24 @@ def f(x):
         inp = torch.randn(6, 7)
         self.assertEqual(gm(inp), f(inp))
 
+    def test_pre_autograd_simple(self):
+        def f(x):
+            y = torch.ones_like(x)
+            return torch.matmul(x, y)
+
+        gm, _ = torch._dynamo.export(
+            f, torch.randn(5, 5), aten_graph=True, pre_autograd=True, tracing_mode="fake"
+        )
+
+        inp = torch.randn(6, 6)
+        self.assertEqual(gm(inp), f(inp))
+        self.assertExpectedInline(gm.code.strip(), """\
+def forward(self, x):
+    arg0, = fx_pytree.tree_flatten_spec(([x], {}), self._in_spec)
+    ones_like_default = torch.ops.aten.ones_like.default(arg0, pin_memory = False)
+    matmul_default = torch.ops.aten.matmul.default(arg0, ones_like_default);  arg0 = ones_like_default = None
+    return pytree.tree_unflatten([matmul_default], self._out_spec)""")
+
     @patch.object(torch._dynamo.config, "dynamic_shapes", True)
     @patch.object(torch._dynamo.config, "capture_scalar_outputs", True)
     def test_export_cond_in_aten_symbolic(self):

torch/_dynamo/eval_frame.py

@@ -605,6 +605,7 @@ def export(
     f: Callable[..., Any],
     *args,
     aten_graph: bool = False,
+    pre_autograd: bool = False,
     decomposition_table: Optional[
         Dict[torch._ops.OpOverload, Callable[..., Any]]
     ] = None,
@@ -623,6 +624,12 @@ def export(
         aten_graph (bool): If True, exports a graph with ATen operators.
         If False, exports a graph with Python operators. Default is False.
 
+        pre_autograd (bool): If True, exports a graph with ATen operators,
+        but before autograd has run. This can be useful if you want to apply further tranformations
+        on a graph before running it through autograd.
+        This flag is only valid if aten_graph=True is set.
+        Default is False.
+
         decomposition_table (dict): A dictionary that maps operators to their decomposition functions.
         Required if aten_graph or tracing_mode is specified. Default is None.
 
@@ -650,6 +657,8 @@ def export(
         assert (
             aten_graph
         ), "Specifying a decomposition_table table or tracing mode is illegal without setting aten_graph=True"
+    if pre_autograd:
+        assert aten_graph, 'pre_autograd=True can only be used when aten_graph=True'
     f = innermost_fn(f)
 
     graph = None
@@ -801,6 +810,7 @@ def graph_with_interpreter(*args):
                 graph_with_interpreter,
                 decomposition_table=decomposition_table,
                 tracing_mode="real",
+                pre_autograd=pre_autograd,
                 _allow_non_fake_inputs=True,
             )(*example_fake_inputs)
 

torch/fx/experimental/proxy_tensor.py

@@ -414,7 +414,8 @@ def can_handle_tensor(x):
     else:
         constant = None
 
-    track_tensor_tree(out, proxy_out, constant=constant, tracer=tracer)
+    with inside_mode(proxy_mode):
+        track_tensor_tree(out, proxy_out, constant=constant, tracer=tracer)
     return out