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.

[ydwu4]

Closes as issue fixed. Repro: patch ed's diff to test/dynamo/test_export.py:test_export_with_module_layer. It produces the following graph

class GraphModule(torch.nn.Module):
    def forward(self, pred, x):
        arg0: "b8[]"; arg1: "f32[3, 3]"; 
    
        arg0, arg1, = fx_pytree.tree_flatten_spec(([pred, x], {}), self._in_spec)
        l_pred_ = arg0
        l_x_ = arg1
        
        # File: /home/yidi/local/pytorch/test/dynamo/test_export.py:1535, code: y = x * 2
        y = l_x_ * 2
        
        # File: /home/yidi/local/pytorch/torch/nn/modules/linear.py:116, code: return F.linear(input, self.weight, self.bias)
        l__self___linear_weight = self.L__self___linear_weight
        l__self___linear_bias = self.L__self___linear_bias
        
        # File: /home/yidi/local/pytorch/torch/_higher_order_ops/cond.py:116, code: return cond_op(pred, true_fn, false_fn, operands)
        cond_true_0 = self.cond_true_0
        cond_false_0 = self.cond_false_0
        cond = torch.ops.higher_order.cond(l_pred_, cond_true_0, cond_false_0, [l__self___linear_bias, l__self___linear_weight, l_x_, y]);  l_pred_ = cond_true_0 = cond_false_0 = l__self___linear_bias = l__self___linear_weight = l_x_ = y = None
        getitem = cond[0];  cond = None
        return pytree.tree_unflatten([getitem], self._out_spec)
        
    class GraphModule(torch.nn.Module):
        def forward(self, l__self___linear_bias, l__self___linear_weight, l_x_, y_true_branch):
            l__self___linear_bias_1 = l__self___linear_bias
            l__self___linear_weight_1 = l__self___linear_weight
            l_x__1 = l_x_
            
            # File: /home/yidi/local/pytorch/torch/nn/modules/linear.py:116, code: return F.linear(input, self.weight, self.bias)
            linear = torch._C._nn.linear(l_x__1, l__self___linear_weight_1, l__self___linear_bias_1);  l_x__1 = l__self___linear_weight_1 = l__self___linear_bias_1 = None
            
            # File: /home/yidi/local/pytorch/test/dynamo/test_export.py:1537, code: return self.linear(val) * torch.tensor(2) * y
            tensor = torch.tensor(2)
            mul = linear * tensor;  linear = tensor = None
            mul_1 = mul * y_true_branch;  mul = y_true_branch = None
            return (mul_1,)
            
    class GraphModule(torch.nn.Module):
        def forward(self, l__self___linear_bias_1, l__self___linear_weight_1, l_x_, y_true_branch):
            l__self___linear_bias_2 = l__self___linear_bias_1
            l__self___linear_weight_2 = l__self___linear_weight_1
            l_x__1 = l_x_
            
            # File: /home/yidi/local/pytorch/torch/nn/modules/linear.py:116, code: return F.linear(input, self.weight, self.bias)
            linear = torch._C._nn.linear(l_x__1, l__self___linear_weight_2, l__self___linear_bias_2);  l_x__1 = l__self___linear_weight_2 = l__self___linear_bias_2 = None
            
            # File: /home/yidi/local/pytorch/test/dynamo/test_export.py:1540, code: return self.linear(val) * torch.tensor(-1)
            tensor = torch.tensor(-1)
            mul = linear * tensor;  linear = tensor = None
            return (mul,)