mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
47dbfecd37
This PR:
- adds a new torch.library.register_fake and deprecates
torch.library.impl_abstract. The motivation is that we have a lot of
confusion around the naming so we are going to align the naming with
the actual subsystem (FakeTensor).
- renames `m.impl_abstract_pystub("fbgemm_gpu.sparse_ops")` to
`m.has_python_registration("fbgemm_gpu.sparse_ops")`. No deprecation
here yet; I need to test how this works with static initialization.
- Renames a bunch of internals to match (e.g. abstractimplpystub ->
pystub)
I'm scared to rename the Python-side internal APIs (e.g.
torch._library.abstract_impl) because of torch.package concerns. I'll do
that in its own isolated PR next just in case it causes problems.
DEPRECATION NOTE: torch.library.impl_abstract was renamed to to
torch.library.register_fake. Please use register_fake. We'll delete
impl_abstract in a future version of PyTorch.
Test Plan:
- existing tests
Pull Request resolved: https://github.com/pytorch/pytorch/pull/123937
Approved by: https://github.com/albanD
43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
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)
|