-
Notifications
You must be signed in to change notification settings - Fork 22k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
New example for preserve_node_meta #128681
base: main
Are you sure you want to change the base?
Conversation
Docstring updated with a new example as suggested by @ezyang and @angelayi Fixes pytorch#127915
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/128681
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit b9613f5 with merge base eff01bc ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
>>> for node in fxg.graph.nodes: | ||
>>> print(node, node.meta) | ||
|
||
>>> def simple_transform(g: torch.fx.Graph) -> torch.fx.Graph: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
hmm, this is not really the transformation I was referring to. It should be with fx.Transformer
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hi @angelayi, a quick clarification, please see the code below. When I do a node traversal with run_node
, it does preserve the meta data.
However, if I do SimpleTransform(fxg).transform()
, it overwrites the existing meta data. I understand transform()
returns a new graph but is this the expected behavior of preserve_node_meta
when called against transform()
?
(tested with Pytorch built from commit eff01bc)
class SimpleAdd(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x, y):
return x + y
class SimpleTransform(torch.fx.Transformer):
def __init__(self, module):
super().__init__(module)
def run_node(self, n: torch.fx.Node):
n.meta["node_op"] = n.op
return super().run_node(n)
model = SimpleAdd()
with torch.fx.traceback.preserve_node_meta():
fxg = torch.fx.symbolic_trace(model)
# node meta prior to transform
"""
x {'seq_nr': -1}
y {'seq_nr': -1}
add {'seq_nr': -1}
output {'seq_nr': -1}
"""
for node in fxg.graph.nodes:
print(node, node.meta)
# node traversal to add new meta
SimpleTransform(fxg).run()
"""
x {'seq_nr': -1, 'node_op': 'placeholder'}
y {'seq_nr': -1, 'node_op': 'placeholder'}
add {'seq_nr': -1, 'node_op': 'call_function'}
output {'seq_nr': -1, 'node_op': 'output'}
"""
for node in fxg.graph.nodes:
print(node, node.meta)
# call super().transform()
fxgn = SimpleTransform(fxg).transform()
# node meta after transform
"""
x {}
y {}
add {'from_node': [('add', <built-in function add>)], 'seq_nr': -1}
output {}
"""
for node in fxgn.graph.nodes:
print(node, node.meta)
Thanks!
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
Docstring updated with a new example as suggested by @ezyang and @angelayi
Fixes #127915
cc @svekars @brycebortree @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv