Fix constant_fold deleting attrs still referenced by live get_attr nodes

[amanyagami]

Summary

torchao.quantization.pt2e.constant_fold.constant_fold's cleanup pass deleted the underlying module attribute for a get_attr node as soon as it found one dead get_attr node for that target — even if another get_attr node with the same target was still live elsewhere in the graph. That leaves the graph with a live get_attr node pointing at an attribute that no longer exists, which graph.lint() (called from eliminate_dead_code) rejects:

RuntimeError: Node shared_1 target shared references nonexistent attribute shared of ...

This was hit in ExecuTorch on a model where the same buffer/parameter is referenced through multiple get_attr nodes.

Fix

Compute the set of targets that still have at least one live get_attr node before deleting anything, and only delattr a target when it isn't in that set — i.e. only once every get_attr node for that target is dead.

Test

Added test/quantization/pt2e/test_constant_fold.py using the exact repro from the issue: two get_attr nodes sharing one buffer target, one folded away (dead), one live.

• Verified it fails with the issue's exact RuntimeError on the unpatched code.
• Passes with the fix, and the resulting graph module still runs and produces correct values for both outputs.
• The pre-existing test/quantization/pt2e/test_quantize_pt2e.py suite (70 tests, which exercises constant_fold indirectly via lowering/quantize_pt2e) still passes.

Fixes #4420

🤖 Generated with Claude Code

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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