pytorch / pytorch

Tensors and Dynamic neural networks in Python with strong GPU acceleration
https://pytorch.org
Other
80.63k stars 21.65k forks source link

SymNode.__repr__ shouldn't be used for compilation output #129403

Closed aorenste closed 1 week ago

aorenste commented 3 weeks ago

🐛 Describe the bug

SymNode.repr only prints out the underlying sympy.Expr without any of the "extra" information that it contains. This results in things like s0*2 != s0*2 because the shape_env is different. Unfortunately this can't easily be changed because repr is used by GraphModule.recompile() to create a pythonic version of the graph. Instead we should have a separate printing function that we use for producing code which would allow us to add more debugging info to repr.

Versions

PyTorch version: 2.5.0a0+git64abc3f

cc @ezyang @anijain2305 @chauhang @penguinwu

bt2513 commented 3 weeks ago

Hi! If no one is currently working on this issue, I would like to take it as a first contribution

aorenste commented 3 weeks ago

Hi! If no one is currently working on this issue, I would like to take it as a first contribution

You're certainly welcome to do it and submit a PR - but since this issue hasn't been triaged yet it's possible that it will be rejected. It would be a good learning experience either way.

ezyang commented 3 weeks ago

Having a separate debugging printout seems good, but defaulting to printing s0+2 seems important because otherwise it would be impossibly wordy

aorenste commented 3 weeks ago

Having a separate debugging printout seems good, but defaulting to printing s0+2 seems important because otherwise it would be impossibly wordy

When you say "defaulting to printing" do you mean __str__ or __repr__ (or both)? My understanding was that in general __repr__ is supposed to be somewhat wordy because it should give you enough information to recreate the object.

eellison commented 3 weeks ago

I was also confused about symints and symnodes having the same print out when working on them.

bt2513 commented 3 weeks ago

awesome, happy to work on that then!

ezyang commented 3 weeks ago

OK, I guess I have less objection to SymNode having some wordier repr (I care a lot about SymInt having a repr like "s0+1" though). However, @aorenste, in general this is a complicated data structure with pointers, you're never going to have a repr that exactly replicates the semantics.

bt2513 commented 3 weeks ago

I can work on this as my first contribution this weekend and create a PR. Could I please be assigned to this task? And who should I put as reviewers? @aorenste @ezyang Thank you!

aorenste commented 3 weeks ago

I can work on this as my first contribution this weekend and create a PR. Could I please be assigned to this task? And who should I put as reviewers? @aorenste @ezyang Thank you!

Sure - put me as a reviewer.