New example for preserve_node_meta

[TharinduRusira]

Docstring updated with a new example as suggested by @ezyang and @angelayi

Fixes #127915

cc @svekars @brycebortree @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

✅ 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

• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 5, 5, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCUDA::test_comprehensive_nn_functional_avg_pool3d_cuda_float16

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

[angelayi]

hmm, this is not really the transformation I was referring to. It should be with fx.Transformer.

[TharinduRusira]

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!

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.