frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[dynamo] Support fx.traceback.annotate as decorator (#165805)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165805
Approved by: https://github.com/Lucaskabela, https://github.com/SherlockNoMad, https://github.com/yushangdi
This commit is contained in:
Animesh Jain
2025-10-17 16:38:12 -07:00
committed by PyTorch MergeBot
parent 23417ae50f
commit d9f94e0d7d
2 changed files with 55 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
torch/_dynamo/variables/torch.py
View File

@ -126,6 +126,7 @@ supported_ctx_manager_classes = dict.fromkeys(
torch.cpu.amp.autocast_mode.autocast,
torch.cuda.amp.autocast_mode.autocast,
torch.fx.traceback.annotate,
torch.fx.traceback.annotate.__wrapped__, # type: ignore[attr-defined]
# We'll let Dynamo inline into the contextlib part of these context
# manager instances, all the way till it invokes the wrapped function
# itself (at which point we wrap it back to special context manager
@ -364,7 +365,10 @@ class TorchCtxManagerClassVariable(BaseTorchVariable):
assert len(args) <= 1 and len(kwargs) == 0
inf_mode = args[0].as_python_constant() if len(args) == 1 else True
return InferenceModeVariable.create(tx, inf_mode)
elif self.value is torch.fx.traceback.annotate:
elif self.value in (
torch.fx.traceback.annotate,
torch.fx.traceback.annotate.__wrapped__, # type: ignore[attr-defined]
):
assert len(args) <= 1 and len(kwargs) == 0
return FxTracebackAnnotateVariable(
args[0].as_python_constant(), source=self.source