From 4ba51bb2c452f953940088cff21ca62bd43e540f Mon Sep 17 00:00:00 2001
From: drisspg <drisspguessous@gmail.com>
Date: Thu, 21 Mar 2024 22:16:49 +0000
Subject: [PATCH] Add keys used for templated attention impls (#122423)

# 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
---
 torchgen/model.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/torchgen/model.py b/torchgen/model.py
index 1f95d34d94d0b..a3f4d9e109888 100644
--- a/torchgen/model.py
+++ b/torchgen/model.py
@@ -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()