File size: 1,413 Bytes
d1ceb73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
# mypy: allow-untyped-defs
from .abstract_impl import AbstractImplHolder
__all__ = ["SimpleLibraryRegistry", "SimpleOperatorEntry", "singleton"]
class SimpleLibraryRegistry:
"""Registry for the "simple" torch.library APIs
The "simple" torch.library APIs are a higher-level API on top of the
raw PyTorch DispatchKey registration APIs that includes:
- fake impl
Registrations for these APIs do not go into the PyTorch dispatcher's
table because they may not directly involve a DispatchKey. For example,
the fake impl is a Python function that gets invoked by FakeTensor.
Instead, we manage them here.
SimpleLibraryRegistry is a mapping from a fully qualified operator name
(including the overload) to SimpleOperatorEntry.
"""
def __init__(self):
self._data = {}
def find(self, qualname: str) -> "SimpleOperatorEntry":
if qualname not in self._data:
self._data[qualname] = SimpleOperatorEntry(qualname)
return self._data[qualname]
singleton: SimpleLibraryRegistry = SimpleLibraryRegistry()
class SimpleOperatorEntry:
"""This is 1:1 to an operator overload.
The fields of SimpleOperatorEntry are Holders where kernels can be
registered to.
"""
def __init__(self, qualname: str):
self.qualname: str = qualname
self.abstract_impl: AbstractImplHolder = AbstractImplHolder(qualname)
|