File size: 1,437 Bytes
c61ccee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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:

    - abstract 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 abstract 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)