Update on "[symbolic shapes] Log SymNode id for provenance"

angelayi · angelayi · commit 92c0f4aef33b · 2025-02-11T14:15:37.000-08:00
We can use the SymNode id to point us back to how previous expressions were created, and construct this nice tree in tlparse: <img width="761" alt="image" src="https://github.com/user-attachments/assets/531b03e8-4398-4d0a-bd11-16078256041c" /> cc ezyang SherlockNoMad EikanWang jgong5 wenzhe-nrv [ghstack-poisoned]
diff --git a/torch/fx/experimental/sym_node.py b/torch/fx/experimental/sym_node.py
@@ -1295,9 +1295,9 @@ def wrapper(self, other=None):
                     del frame
 
                 if other is not None:
-                    arguments = [str(self), str(other)]
+                    arguments = [self, other]
                 else:
-                    arguments = [str(self)]
+                    arguments = [self]
 
                 def get_id(sym_node) -> Optional[int]:
                     # We don't want to return an ID if the input is a constant
@@ -1309,8 +1309,10 @@ def get_id(sym_node) -> Optional[int]:
                         "method": method,
                         "result": str(result),
                         "result_id": id(result),
-                        "arguments": arguments,
-                        "argument_ids": [i for i in arguments if i is not None],
+                        "arguments": [str(a) for a in arguments],
+                        "argument_ids": [
+                            get_id(i) for i in arguments if get_id(i) is not None
+                        ],
                         "user_bottom_stack": str(user_bottom_stack),
                         "user_top_stack": str(user_top_stack),
                         "floc": str(floc),
