Dynamo and cond with free variables creates malformed graph

[ezyang]

🐛 Describe the bug

repro

diff --git a/test/dynamo/test_export.py b/test/dynamo/test_export.py
index b0640f6511..8fb8754c51 100644
--- a/test/dynamo/test_export.py
+++ b/test/dynamo/test_export.py
@@ -1443,8 +1443,10 @@ class ExportTests(torch._dynamo.test_case.TestCase):
                 self.linear = torch.nn.Linear(3, 3)
 
             def forward(self, pred, x):
+                y = x * 2
+
                 def true_fn(val):
-                    return self.linear(val) * torch.tensor(2)
+                    return self.linear(val) * torch.tensor(2) * y
 
                 def false_fn(val):
                     return self.linear(val) * torch.tensor(-1)

the true graph ends up being

def forward(self, x):
    self_linear = self.self_linear(x);  x = None
    tensor = torch.tensor(2)
    mul = self_linear * tensor;  self_linear = tensor = None
    mul_1 = mul * mul;  mul = mul = None
    return mul_1

which is so bad

Versions

master

cc @soumith @msaroufim @wconstab @ngimel @bdhirsh @mlazos @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @chunyuan-w @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

[ezyang]

cc @voznesenskym

[avikchaudhuri]

I played around with the repro a bit.

def forward(self, pred, x):
    y = x + 2 // note: add instead of mul

    def true_fn(val):
        return self.linear(val) * torch.tensor(2) * y
 
    def false_fn(val):
        return self.linear(val) * torch.tensor(-1) * y

With this variant the generated code actually crashes because it tries to look up a variable add which is not in scope (but it should be, because this is what the addition stored in y gets renamed to):

def forward(self, x):
    self_linear = self.self_linear(x);  x = None
    tensor = torch.tensor(2)
    mul = self_linear * tensor;  self_linear = tensor = None
    mul_1 = mul * add;  mul = mul = None
    return mul_1

Basically this means that the mul * mul in the original repro is just a coincidence.