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

Add keys used for templated attention impls (#122423)

Browse Source 
# Summary

Mypy will complain that these attributes dont exist for this PR: https://github.com/pytorch/pytorch/pull/121845/
Pull Request resolved: https://github.com/pytorch/pytorch/pull/122423
Approved by: https://github.com/bdhirsh
This commit is contained in:
drisspg
2024-03-21 22:16:49 +00:00
committed by PyTorch MergeBot
parent 224beecee6
commit 4ba51bb2c4
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
torchgen/model.py
View File

@ -93,7 +93,9 @@ class DispatchKey(Enum):
NestedTensor = auto()
Dense = auto()
PythonTLSSnapshot = auto()
PreDispatch = auto()
PythonDispatcher = auto()
Python = auto()
FuncTorchDynamicLayerBackMode = auto()
ZeroTensor = auto()
@ -106,6 +108,8 @@ class DispatchKey(Enum):
AutogradNestedTensor = auto()
Tracer = auto()
Autocast = auto()
AutocastCPU = auto()
AutocastCUDA = auto()
Batched = auto()
VmapMode = auto()
FuncTorchGradWrapper = auto()