Repository: pytorch/torchdynamo

Issue:
`cond` won't compile if one branch returns its argument
Repro:
```
import torch
import torchdynamo

def true_fn(x):
    return x

def false_fn(x):
    return x + x

def f(pred, x):
    return cond(pred, true_fn, false_fn, [x])

opt_fn = torchdynamo.optimize("eager")(f)
opt_fn(torch.tensor(True), torch.ones(2, 3))
```

Compilation of the true branch fails with:
```
AssertionError: whole graph export entails exactly one call
```

cc @voznesenskym 


Discussion:
This is a real bug - cond relies on export, and I bet it's failing to graph capture true_fn. Good find, thank you, I'll fix it ASAP. 

Resolving patch:
diff --git a/torchdynamo/convert_frame.py b/torchdynamo/convert_frame.py
--- a/torchdynamo/convert_frame.py
+++ b/torchdynamo/convert_frame.py
@@ -215,7 +215,9 @@ def format_error_msg(exc, frame):
     return msg
 
 
-def convert_frame_assert(compiler_fn: Callable, guard_export_fn=None, one_graph=True):
+def convert_frame_assert(
+    compiler_fn: Callable, guard_export_fn=None, one_graph=True, export=False
+):
     """Fully convert a frame into an FX graph"""
     init_logging()
 
@@ -294,6 +296,7 @@ def transform(instructions, code_options):
                 code_options,
                 compiler_fn,
                 one_graph,
+                export,
             )
             tracer.run()
             output = tracer.output
diff --git a/torchdynamo/eval_frame.py b/torchdynamo/eval_frame.py
--- a/torchdynamo/eval_frame.py
+++ b/torchdynamo/eval_frame.py
@@ -474,7 +474,9 @@ def result_capturing_wrapper(*graph_inputs):
     remove_from_cache(f)
     with patch(f"{__name__}.most_recent_backend", None):
         opt_f = optimize_assert(
-            dynamo_normalization_capturing_compiler, guard_export_fn=guard_export_print
+            dynamo_normalization_capturing_compiler,
+            guard_export_fn=guard_export_print,
+            export=True,
         )(f)
         # TODO(voz): We may have instances of `f` that mutate inputs, we should track sideffects and reject.
         result_traced = opt_f(*args, **kwargs)
@@ -531,7 +533,7 @@ def graph_with_interpreter(*args):
     return (new_graph, out_guards)
 
 
-def optimize_assert(backend, *, guard_export_fn=None):
+def optimize_assert(backend, *, guard_export_fn=None, export=False):
     """
     The same as `torchdynamo.optimize(backend, nopython=True)`
     """
@@ -541,7 +543,8 @@ def optimize_assert(backend, *, guard_export_fn=None):
     backend_ctx_ctor = getattr(backend, "backend_ctx_ctor", null_context)
 
     return _optimize_catch_errors(
-        convert_frame.convert_frame_assert(backend, guard_export_fn), backend_ctx_ctor
+        convert_frame.convert_frame_assert(backend, guard_export_fn, export=export),
+        backend_ctx_ctor,
     )
 
 
diff --git a/torchdynamo/symbolic_convert.py b/torchdynamo/symbolic_convert.py
--- a/torchdynamo/symbolic_convert.py
+++ b/torchdynamo/symbolic_convert.py
@@ -1291,6 +1291,7 @@ def __init__(
         code_options,
         compiler_fn,
         one_graph,
+        export,
     ):
         super(InstructionTranslator, self).__init__(
             output=OutputGraph(f_globals, code_options, compiler_fn, self),
@@ -1304,6 +1305,12 @@ def __init__(
             f_code=f_code,
         )
         self.one_graph: bool = one_graph
+        self.export = export
+        if self.export:
+            assert (
+                self.one_graph
+            ), "Export without one graph - something has gone wrong."
+
         vars = list(code_options["co_varnames"])
         vars.extend(x for x in self.cell_and_freevars() if x not in vars)
         self.symbolic_locals = collections.OrderedDict(
@@ -1408,7 +1415,7 @@ def create_call_resume_at(self, inst):
         return cg.get_instructions()
 
     def RETURN_VALUE(self, inst):
-        if self.output.count_calls() == 0:
+        if self.output.count_calls() == 0 and not self.export:
             raise exc.SkipFrame()
         self.instruction_pointer = None
         self.output.compile_subgraph(self)


Test patch:
diff --git a/test/test_misc.py b/test/test_misc.py
--- a/test/test_misc.py
+++ b/test/test_misc.py
@@ -2280,6 +2280,30 @@ def f(pred, pred2, x):
             same(torch.tensor([0.0, 0.0]), false_false_sum_neg)
         )  # * -1 then add x
 
+    @patch.object(torchdynamo.config, "fake_tensor_propagation", False)
+    def test_cond_export_single_arg(self):
+        from functorch.experimental.cond import cond
+
+        def true_fn(x):
+            return x
+
+        def false_fn(x):
+            return x.sin()
+
+        def f(pred, x):
+            return cond(pred, true_fn, false_fn, [x])
+
+        graph, guard = torchdynamo.export(
+            f, torch.tensor(False), torch.tensor([0.25, 0.25])
+        )
+        true_mirror = graph(torch.tensor(True), torch.tensor([0.25, 0.25]))
+        self.assertTrue(same(torch.tensor([0.25, 0.25]), true_mirror))
+        true_mirror_2 = graph(torch.tensor(True), torch.tensor([0.33, 0.33, 0.33]))
+        self.assertTrue(same(torch.tensor([0.33, 0.33, 0.33]), true_mirror_2))
+
+        false_sin = graph(torch.tensor(False), torch.tensor([0.5, 0.5]))
+        self.assertTrue(same(torch.sin(torch.tensor([0.5, 0.5])), false_sin))
+
     def test_disable_optimize(self):
         cnt = torchdynamo.testing.CompileCounter()