Lines Matching refs:dispatch_key

264         for dispatch_key in _HIGHER_ORDER_OP_DEFAULT_FALLTHROUGH_DISPATCH_KEYS:
265             self.fallthrough(dispatch_key)
284     def fallthrough(self, dispatch_key):  argument
285         self.non_fallthrough_keys = self.non_fallthrough_keys.remove(dispatch_key)
289     def dispatch(self, /, dispatch_key, *args, **kwargs):  argument
292         if dispatch_key in self._dispatch_cache:
293             kernel = self._dispatch_cache[dispatch_key]
297         if dispatch_key == DispatchKey.FuncTorchDynamicLayerFrontMode:
300         if dispatch_key == DispatchKey.Python:
370 …     functionality_key = torch._C._to_functionality_key(dispatch_key)  # type: ignore[attr-defined]
392         final_key = resolve_key(self, dispatch_key)
406         if dispatch_key != DispatchKey.PreDispatch:
407             self._dispatch_cache[dispatch_key] = self.py_kernels[final_key]
961         dispatch_key = dispatch_key_set.highestPriorityTypeId()
964             self._get_dispatch(dispatch_key)
965             if dispatch_key not in self._dispatch_cache
966             else self._dispatch_cache[dispatch_key]
973                 self.name(), dispatch_key
976                     args, kwargs, fallthrough_keys + [dispatch_key]