Update on "Rewrite fallthrough to more closely match how C++ works"

ezyang · ezyang · commit f7a00006682c · 2023-03-08T11:16:38.000-08:00
Fallthrough is modeled as a mask which we use to remove keys from the compute dispatch key set for eligibility. It's possible this addresses #89037 in a better way than #95891 but I cannot easily tell as the original repro no longer works and the new PR does not have a test. Signed-off-by: Edward Z. Yang <ezyangmeta.com> [ghstack-poisoned]
diff --git a/torch/_C/__init__.pyi.in b/torch/_C/__init__.pyi.in
@@ -1325,12 +1325,15 @@ class DispatchKeySet:
     def __and__(self, other: DispatchKeySet) -> DispatchKeySet: ...
     def highestPriorityTypeId(self) -> DispatchKey: ...
     def has(self, k: _dispatchkey) -> _bool: ...
+    def add(self, k: _dispatchkey) -> DispatchKeySet: ...
+    def remove(self, k: _dispatchkey) -> DispatchKeySet: ...
     def __repr__(self) -> str: ...
 
 _dispatch_autogradother_backends: DispatchKeySet
 
 def _dispatch_has_backend_fallback(dispatch: _dispatchkey) -> _bool: ...
 def _dispatch_keyset_full_after(t: _dispatchkey) -> DispatchKeySet: ...
+def _dispatch_keyset_full() -> DispatchKeySet: ...
 def _dispatch_keyset_to_string(keyset: DispatchKeySet) -> str: ...
 def _dispatch_get_backend_keyset_from_autograd(
     dispatch: _dispatchkey,
diff --git a/torch/_ops.py b/torch/_ops.py
@@ -3,7 +3,7 @@
 import inspect
 import sys
 import types
-from typing import Any, Dict, Type
+from typing import Any, Callable, Dict, Type, Union
 
 import torch._C
 
@@ -82,7 +82,9 @@ def __init__(self):
         # thought of as an open world extension of dispatch keys, so it
         # makes sense that you should be able to register them, the same
         # way you can register dispatch keys.
-        self.python_key_mode_table: Dict[Type[TorchDispatchMode]] = {}
+        self.python_key_mode_table: Dict[
+            Type[TorchDispatchMode], Callable[..., Any]
+        ] = {}
 
         # This table allows you to override the behavior of functorch
         # transformations.  NB: this currently only does something for
@@ -115,7 +117,7 @@ def inner(fn):
 
             if k in self.py_kernels:
                 raise RuntimeError(
-                    f"Trying to override a python impl for {dispatch_key_or_mode} on operator {self._name}"
+                    f"Trying to override a python impl for {k} on operator {self.name()}"
                 )
             self.py_kernels[k] = fn
             self._dispatch_cache.clear()
