Fix AttrProxy module reuse in Dynamo (#185300)
Dynamo checks side-effect tracking before VariableBuilder wraps a value, but proxy_tensor _AttrProxy modules are unwrapped inside wrap_module before they are tracked. When the same proxy-backed module is first seen as a local key and later through WeakKeyDictionary internals, the early lookup checks the proxy id, misses the real module that was already tracked, unwraps to the same base module, and then trips the duplicate mutation-tracking assertion. Fix this by re-checking side-effect tracking immediately after _AttrProxy is unwrapped to its base module. The reuse path mirrors the existing duplicate-object handling: install a duplicate guard when possible and refresh the UnspecializedNNModuleVariable stack source before returning the existing tracker. This keeps the fix scoped to the proxy-module path instead of relaxing the global mutation-tracking assertion. Test Plan: python test/dynamo/test_dicts.py -k test_weakkeydict_attr_proxy_key python test/dynamo/test_dicts.py -k weakref_dict lintrunner -a Benchmark Results: Duplicate AttrProxy WeakKeyDictionary compile path: main failed with AssertionError after 3545.040 ms; fix-161053 passed in 1982.965 ms. Non-duplicate AttrProxy compile microbenchmark over 50 iterations: main median 15.008 ms, mean 15.084 ms; fix median 15.467 ms, mean 15.574 ms. Fixes #161053 Generated by my agent Pull Request resolved: https://github.com/pytorch/pytorch/pull/185300 Approved by: https://github.com/bobrenjc93
J
Jason Ansel committed
3b0b5dc065ff58e120040014a08d998d88c6eed1
Parent: 7f098f9
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
on 8/24/2026, 4:51:58 AM