update test files

Author: youkaichao

tests/depyf_output/debug_function_aot_eager_with_dynamic_shape_with_grad/__compiled_fn_1.Backward_graph.0.py

@@ -1,20 +1,19 @@
 from __future__ import annotations
-
-
-
-def forward(self, primals_1: "f32[10]", tangents_1: "f32[10]"):
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    neg: "f32[10]" = torch.ops.aten.neg.default(tangents_1)
-    abs_1: "f32[10]" = torch.ops.aten.abs.default(primals_1)
-    add: "f32[10]" = torch.ops.aten.add.Tensor(abs_1, 1);  abs_1 = None
-    div_1: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
-    div_2: "f32[10]" = torch.ops.aten.div.Tensor(div_1, add);  div_1 = None
-    mul: "f32[10]" = torch.ops.aten.mul.Tensor(neg, div_2);  neg = div_2 = None
-    div_3: "f32[10]" = torch.ops.aten.div.Tensor(tangents_1, add);  tangents_1 = add = None
-    sgn: "f32[10]" = torch.ops.aten.sgn.default(primals_1);  primals_1 = None
-    mul_1: "f32[10]" = torch.ops.aten.mul.Tensor(mul, sgn);  mul = sgn = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    add_1: "f32[10]" = torch.ops.aten.add.Tensor(div_3, mul_1);  div_3 = mul_1 = None
-    return (add_1, None)
-    
+import torch
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "f32[10]", tangents_1: "f32[10]"):
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        neg: "f32[10]" = torch.ops.aten.neg.default(tangents_1)
+        abs_1: "f32[10]" = torch.ops.aten.abs.default(primals_1)
+        add: "f32[10]" = torch.ops.aten.add.Tensor(abs_1, 1);  abs_1 = None
+        div_1: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
+        div_2: "f32[10]" = torch.ops.aten.div.Tensor(div_1, add);  div_1 = None
+        mul: "f32[10]" = torch.ops.aten.mul.Tensor(neg, div_2);  neg = div_2 = None
+        div_3: "f32[10]" = torch.ops.aten.div.Tensor(tangents_1, add);  tangents_1 = add = None
+        sgn: "f32[10]" = torch.ops.aten.sgn.default(primals_1);  primals_1 = None
+        mul_1: "f32[10]" = torch.ops.aten.mul.Tensor(mul, sgn);  mul = sgn = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        add_1: "f32[10]" = torch.ops.aten.add.Tensor(div_3, mul_1);  div_3 = mul_1 = None
+        return (add_1, None)
+        

tests/depyf_output/debug_function_aot_eager_with_dynamic_shape_with_grad/__compiled_fn_1.Captured_Graph.0.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
-
-
-
-def forward(self, L_a_: "f32[10]", L_b_: "f32[10]"):
-    l_a_ = L_a_
-    l_b_ = L_b_
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    abs_1: "f32[10]" = torch.abs(l_a_)
-    add: "f32[10]" = abs_1 + 1;  abs_1 = None
-    x: "f32[10]" = l_a_ / add;  l_a_ = add = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
-    sum_1: "f32[]" = l_b_.sum();  l_b_ = None
-    lt: "b8[]" = sum_1 < 0;  sum_1 = None
-    return (x, lt)
-    
+import torch
+class GraphModule(torch.nn.Module):
+    def forward(self, L_a_: "f32[10]", L_b_: "f32[10]"):
+        l_a_ = L_a_
+        l_b_ = L_b_
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        abs_1: "f32[10]" = torch.abs(l_a_)
+        add: "f32[10]" = abs_1 + 1;  abs_1 = None
+        x: "f32[10]" = l_a_ / add;  l_a_ = add = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
+        sum_1: "f32[]" = l_b_.sum();  l_b_ = None
+        lt: "b8[]" = sum_1 < 0;  sum_1 = None
+        return (lt, x)
+        

tests/depyf_output/debug_function_aot_eager_with_dynamic_shape_with_grad/__compiled_fn_1.Forward_graph.0.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
-
-
-
-def forward(self, primals_1: "f32[10]", primals_2: "f32[10]"):
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    abs_1: "f32[10]" = torch.ops.aten.abs.default(primals_1)
-    add: "f32[10]" = torch.ops.aten.add.Tensor(abs_1, 1);  abs_1 = None
-    div: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add);  add = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
-    sum_1: "f32[]" = torch.ops.aten.sum.default(primals_2);  primals_2 = None
-    lt: "b8[]" = torch.ops.aten.lt.Scalar(sum_1, 0);  sum_1 = None
-    return (div, lt, primals_1)
-    
+import torch
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "f32[10]", primals_2: "f32[10]"):
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        abs_1: "f32[10]" = torch.ops.aten.abs.default(primals_1)
+        add: "f32[10]" = torch.ops.aten.add.Tensor(abs_1, 1);  abs_1 = None
+        div: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add);  add = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
+        sum_1: "f32[]" = torch.ops.aten.sum.default(primals_2);  primals_2 = None
+        lt: "b8[]" = torch.ops.aten.lt.Scalar(sum_1, 0);  sum_1 = None
+        return (lt, div, primals_1)
+        

tests/depyf_output/debug_function_aot_eager_with_dynamic_shape_with_grad/__compiled_fn_1.Joint_graph.0.py

@@ -1,30 +1,29 @@
 from __future__ import annotations
-
-
-
-def forward(self, primals, tangents):
-    primals_1: "f32[10]"; primals_2: "f32[10]"; tangents_1: "f32[10]"; 
-
-    primals_1, primals_2, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    abs_1: "f32[10]" = torch.ops.aten.abs.default(primals_1)
-    add: "f32[10]" = torch.ops.aten.add.Tensor(abs_1, 1);  abs_1 = None
-    div: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
+import torch
+class joint_helper(torch.nn.Module):
+    def forward(self, primals, tangents):
+        primals_1: "f32[10]"; primals_2: "f32[10]"; tangents_1: "f32[10]"; 
 
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
-    sum_1: "f32[]" = torch.ops.aten.sum.default(primals_2);  primals_2 = None
-    lt: "b8[]" = torch.ops.aten.lt.Scalar(sum_1, 0);  sum_1 = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    neg: "f32[10]" = torch.ops.aten.neg.default(tangents_1)
-    div_1: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
-    div_2: "f32[10]" = torch.ops.aten.div.Tensor(div_1, add);  div_1 = None
-    mul: "f32[10]" = torch.ops.aten.mul.Tensor(neg, div_2);  neg = div_2 = None
-    div_3: "f32[10]" = torch.ops.aten.div.Tensor(tangents_1, add);  tangents_1 = add = None
-    sgn: "f32[10]" = torch.ops.aten.sgn.default(primals_1);  primals_1 = None
-    mul_1: "f32[10]" = torch.ops.aten.mul.Tensor(mul, sgn);  mul = sgn = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    add_1: "f32[10]" = torch.ops.aten.add.Tensor(div_3, mul_1);  div_3 = mul_1 = None
-    return pytree.tree_unflatten([div, lt, add_1, None], self._out_spec)
-    
+        primals_1, primals_2, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        abs_1: "f32[10]" = torch.ops.aten.abs.default(primals_1)
+        add: "f32[10]" = torch.ops.aten.add.Tensor(abs_1, 1);  abs_1 = None
+        div: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
+        sum_1: "f32[]" = torch.ops.aten.sum.default(primals_2);  primals_2 = None
+        lt: "b8[]" = torch.ops.aten.lt.Scalar(sum_1, 0);  sum_1 = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        neg: "f32[10]" = torch.ops.aten.neg.default(tangents_1)
+        div_1: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
+        div_2: "f32[10]" = torch.ops.aten.div.Tensor(div_1, add);  div_1 = None
+        mul: "f32[10]" = torch.ops.aten.mul.Tensor(neg, div_2);  neg = div_2 = None
+        div_3: "f32[10]" = torch.ops.aten.div.Tensor(tangents_1, add);  tangents_1 = add = None
+        sgn: "f32[10]" = torch.ops.aten.sgn.default(primals_1);  primals_1 = None
+        mul_1: "f32[10]" = torch.ops.aten.mul.Tensor(mul, sgn);  mul = sgn = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        add_1: "f32[10]" = torch.ops.aten.add.Tensor(div_3, mul_1);  div_3 = mul_1 = None
+        return pytree.tree_unflatten([lt, div, add_1, None], self._out_spec)
+        

tests/depyf_output/debug_function_aot_eager_with_dynamic_shape_with_grad/__compiled_fn_1.pre_insert_deferred_runtime_asserts___compiled_fn_1.0.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
-
-
-
-def forward(self, L_a_: "f32[10]", L_b_: "f32[10]"):
-    l_a_ = L_a_
-    l_b_ = L_b_
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    abs_1: "f32[10]" = torch.abs(l_a_)
-    add: "f32[10]" = abs_1 + 1;  abs_1 = None
-    x: "f32[10]" = l_a_ / add;  l_a_ = add = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
-    sum_1: "f32[]" = l_b_.sum();  l_b_ = None
-    lt: "b8[]" = sum_1 < 0;  sum_1 = None
-    return (x, lt)
-    
+import torch
+class GraphModule(torch.nn.Module):
+    def forward(self, L_a_: "f32[10]", L_b_: "f32[10]"):
+        l_a_ = L_a_
+        l_b_ = L_b_
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        abs_1: "f32[10]" = torch.abs(l_a_)
+        add: "f32[10]" = abs_1 + 1;  abs_1 = None
+        x: "f32[10]" = l_a_ / add;  l_a_ = add = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
+        sum_1: "f32[]" = l_b_.sum();  l_b_ = None
+        lt: "b8[]" = sum_1 < 0;  sum_1 = None
+        return (lt, x)
+        

tests/depyf_output/debug_function_aot_eager_with_dynamic_shape_with_grad/__compiled_fn_1.tensorify_python_scalars.0.py

@@ -1,30 +1,29 @@
 from __future__ import annotations
-
-
-
-def forward(self, primals, tangents):
-    primals_1: "f32[10]"; primals_2: "f32[10]"; tangents_1: "f32[10]"; 
-
-    primals_1, primals_2, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    abs_1: "f32[10]" = torch.ops.aten.abs.default(primals_1)
-    add: "f32[10]" = torch.ops.aten.add.Tensor(abs_1, 1);  abs_1 = None
-    div: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
+import torch
+class joint_helper(torch.nn.Module):
+    def forward(self, primals, tangents):
+        primals_1: "f32[10]"; primals_2: "f32[10]"; tangents_1: "f32[10]"; 
 
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
-    sum_1: "f32[]" = torch.ops.aten.sum.default(primals_2);  primals_2 = None
-    lt: "b8[]" = torch.ops.aten.lt.Scalar(sum_1, 0);  sum_1 = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    neg: "f32[10]" = torch.ops.aten.neg.default(tangents_1)
-    div_1: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
-    div_2: "f32[10]" = torch.ops.aten.div.Tensor(div_1, add);  div_1 = None
-    mul: "f32[10]" = torch.ops.aten.mul.Tensor(neg, div_2);  neg = div_2 = None
-    div_3: "f32[10]" = torch.ops.aten.div.Tensor(tangents_1, add);  tangents_1 = add = None
-    sgn: "f32[10]" = torch.ops.aten.sgn.default(primals_1);  primals_1 = None
-    mul_1: "f32[10]" = torch.ops.aten.mul.Tensor(mul, sgn);  mul = sgn = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
-    add_1: "f32[10]" = torch.ops.aten.add.Tensor(div_3, mul_1);  div_3 = mul_1 = None
-    return pytree.tree_unflatten([div, lt, add_1, None], self._out_spec)
-    
+        primals_1, primals_2, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        abs_1: "f32[10]" = torch.ops.aten.abs.default(primals_1)
+        add: "f32[10]" = torch.ops.aten.add.Tensor(abs_1, 1);  abs_1 = None
+        div: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:5 in toy_function, code: if b.sum() < 0:
+        sum_1: "f32[]" = torch.ops.aten.sum.default(primals_2);  primals_2 = None
+        lt: "b8[]" = torch.ops.aten.lt.Scalar(sum_1, 0);  sum_1 = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        neg: "f32[10]" = torch.ops.aten.neg.default(tangents_1)
+        div_1: "f32[10]" = torch.ops.aten.div.Tensor(primals_1, add)
+        div_2: "f32[10]" = torch.ops.aten.div.Tensor(div_1, add);  div_1 = None
+        mul: "f32[10]" = torch.ops.aten.mul.Tensor(neg, div_2);  neg = div_2 = None
+        div_3: "f32[10]" = torch.ops.aten.div.Tensor(tangents_1, add);  tangents_1 = add = None
+        sgn: "f32[10]" = torch.ops.aten.sgn.default(primals_1);  primals_1 = None
+        mul_1: "f32[10]" = torch.ops.aten.mul.Tensor(mul, sgn);  mul = sgn = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:4 in toy_function, code: x = a / (torch.abs(a) + 1)
+        add_1: "f32[10]" = torch.ops.aten.add.Tensor(div_3, mul_1);  div_3 = mul_1 = None
+        return pytree.tree_unflatten([lt, div, add_1, None], self._out_spec)
+        

tests/depyf_output/debug_function_aot_eager_with_dynamic_shape_with_grad/__compiled_fn_11.Backward_graph.0.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
-
-
-
-def forward(self, primals_1: "Sym(s0)", primals_2: "f32[s0]", primals_3: "f32[s0]", tangents_1: "f32[s0]"):
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:7 in torch_dynamo_resume_in_toy_function_at_5, code: return x * b
-    mul_4: "f32[s0]" = torch.ops.aten.mul.Tensor(tangents_1, primals_3);  primals_3 = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:6 in torch_dynamo_resume_in_toy_function_at_5, code: b = b * -1
-    mul: "f32[s0]" = torch.ops.aten.mul.Tensor(primals_2, -1);  primals_2 = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:7 in torch_dynamo_resume_in_toy_function_at_5, code: return x * b
-    mul_5: "f32[s0]" = torch.ops.aten.mul.Tensor(tangents_1, mul);  tangents_1 = mul = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:6 in torch_dynamo_resume_in_toy_function_at_5, code: b = b * -1
-    mul_6: "f32[s0]" = torch.ops.aten.mul.Tensor(mul_4, -1);  mul_4 = None
-    return (None, mul_6, mul_5)
-    
+import torch
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "Sym(s52)", primals_2: "f32[s52]", primals_3: "f32[s52]", tangents_1: "f32[s52]"):
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:7 in torch_dynamo_resume_in_toy_function_at_5, code: return x * b
+        mul_4: "f32[s52]" = torch.ops.aten.mul.Tensor(tangents_1, primals_3);  primals_3 = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:6 in torch_dynamo_resume_in_toy_function_at_5, code: b = b * -1
+        mul: "f32[s52]" = torch.ops.aten.mul.Tensor(primals_2, -1);  primals_2 = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:7 in torch_dynamo_resume_in_toy_function_at_5, code: return x * b
+        mul_5: "f32[s52]" = torch.ops.aten.mul.Tensor(tangents_1, mul);  tangents_1 = mul = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:6 in torch_dynamo_resume_in_toy_function_at_5, code: b = b * -1
+        mul_6: "f32[s52]" = torch.ops.aten.mul.Tensor(mul_4, -1);  mul_4 = None
+        return (None, mul_6, mul_5)
+        

tests/depyf_output/debug_function_aot_eager_with_dynamic_shape_with_grad/__compiled_fn_11.Captured_Graph.0.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
-
-
-
-def forward(self, s0: "Sym(s0)", L_b_: "f32[s0]", L_x_: "f32[s0]"):
-    l_b_ = L_b_
-    l_x_ = L_x_
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:6 in torch_dynamo_resume_in_toy_function_at_5, code: b = b * -1
-    b: "f32[s0]" = l_b_ * -1;  l_b_ = None
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:7 in torch_dynamo_resume_in_toy_function_at_5, code: return x * b
-    mul_1: "f32[s0]" = l_x_ * b;  l_x_ = b = None
-    return (mul_1,)
-    
+import torch
+class GraphModule(torch.nn.Module):
+    def forward(self, s52: "Sym(s52)", L_b_: "f32[s52]", L_x_: "f32[s52]"):
+        l_b_ = L_b_
+        l_x_ = L_x_
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:6 in torch_dynamo_resume_in_toy_function_at_5, code: b = b * -1
+        b: "f32[s52]" = l_b_ * -1;  l_b_ = None
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:7 in torch_dynamo_resume_in_toy_function_at_5, code: return x * b
+        mul_1: "f32[s52]" = l_x_ * b;  l_x_ = b = None
+        return (mul_1,)
+        

tests/depyf_output/debug_function_aot_eager_with_dynamic_shape_with_grad/__compiled_fn_11.Forward_graph.0.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
-
-
-
-def forward(self, primals_1: "Sym(s0)", primals_2: "f32[s0]", primals_3: "f32[s0]"):
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:6 in torch_dynamo_resume_in_toy_function_at_5, code: b = b * -1
-    mul: "f32[s0]" = torch.ops.aten.mul.Tensor(primals_2, -1)
-    
-     # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:7 in torch_dynamo_resume_in_toy_function_at_5, code: return x * b
-    mul_2: "f32[s0]" = torch.ops.aten.mul.Tensor(primals_3, mul);  mul = None
-    return (mul_2, primals_2, primals_3, primals_1)
-    
+import torch
+class GraphModule(torch.nn.Module):
+    def forward(self, primals_1: "Sym(s52)", primals_2: "f32[s52]", primals_3: "f32[s52]"):
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:6 in torch_dynamo_resume_in_toy_function_at_5, code: b = b * -1
+        mul: "f32[s52]" = torch.ops.aten.mul.Tensor(primals_2, -1)
+        
+         # File: /Users/youkaichao/data/DeepLearning/depyf/tests/test_pytorch/test_pytorch.py:7 in torch_dynamo_resume_in_toy_function_at_5, code: return x * b
+        mul_2: "f32[s52]" = torch.ops.aten.mul.Tensor(primals_3, mul);  mul = None
+        return (mul_2, primals_2, primals_3, primals_1)
+