[FX][CodeGen] Make sure fx code is valid in python

[youkaichao]

This PR fixes two cases when fx generated code is invalid in python (syntax error):

1. multiple type annotation in one line: var1: annotation1, var2: annotation2 = function_call()
2. invalid type annotation for scalars like var1: f32[] = function_call().

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/113345

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 082709a with merge base 8c704f7 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[youkaichao]

For those who are triaging PRs, please route this to @jansel .

[ezyang]

Can you show us the before and after code?

[youkaichao]

Sure.

Before:

the first line primals_1: f32[10], primals_2: f32[10], tangents_1: f32[10], has a syntax error. The line lt: b8[] and sum_1: f32[] also causes syntax error.

def forward(self, primals, tangents):
    primals_1: f32[10], primals_2: f32[10], tangents_1: f32[10], = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
    # File: depyf/tests/test_pytorch/test_debug_function_aot.py:5, 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: depyf/tests/test_pytorch/test_debug_function_aot.py:6, 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: depyf/tests/test_pytorch/test_debug_function_aot.py:5, 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
    sign: f32[10] = torch.ops.aten.sign.default(primals_1);  primals_1 = None
    mul_1: f32[10] = torch.ops.aten.mul.Tensor(mul, sign);  mul = sign = None
    
    # File: depyf/tests/test_pytorch/test_debug_function_aot.py:5, 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)

After:

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: depyf/tests/test_pytorch/test_debug_function_aot.py:5, 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: depyf/tests/test_pytorch/test_debug_function_aot.py:6, 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: depyf/tests/test_pytorch/test_debug_function_aot.py:5, 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
    sign: "f32[10]" = torch.ops.aten.sign.default(primals_1);  primals_1 = None
    mul_1: "f32[10]" = torch.ops.aten.mul.Tensor(mul, sign);  mul = sign = None
    
    # File: depyf/tests/test_pytorch/test_debug_function_aot.py:5, 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)

[ezyang]

This is plausible I guess, although strings aren't valid types. I guess these "type" annotations never were intended to be parsed for real in Python.

[youkaichao]

@ezyang I'm working to make torch.compile debuggable in the https://github.com/thuml/depyf project. I try to link code to fx graphs, and it would be of great help if users can debug these functions. (By debugging, I mean step-by-step executing these functions, say to find which op leads to NaN). Therefore, I need to make these functions valid python code, just like this pr: #111635 .

Plus, since these are meant to be python code, they'd better have correct syntax I think.

[ezyang]

OK, well, you'll have to figure out the test fails first

[youkaichao]

@ezyang I have updated all the failed tests.

[ezyang]

OK

[ezyang]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[Chillee]

@youkaichao This is only from the print_readable output - do you need that to be valid python code? I find the additional strings to be somewhat annoying.

[youkaichao]

@Chillee yes, I use the code here:
https://github.com/thuml/depyf/blob/dbc220a081c6f940071016499f3ac4ef06bc0030/depyf/explain/patched_lazy_format_graph_code.py#L16

The shape annotation greatly helps when users step through the code.

[youkaichao]

@Chillee which aspect makes it annoying?

[Chillee]

@youkaichao It's just additional visual clutter for every line. I would prefer to instead have an option to add it for use cases like yours.

[youkaichao]

My opinion is that readable code should be valid code, too. If someone wants to only read the code, I suppose only adding two quotes "" does not hurt. Human eyes can easily ignore them.

[Chillee]

fwiw, I've since changed my mind on this. Making the readable code valid is quite useful :P