Source code for torch.library

from ._ops import OpOverload
from typing import Set
import traceback
import torch
import weakref

__all__ = ['Library', 'impl', 'define', 'fallthrough_kernel']

# Set containing the combination of (namespace, operator, DispatchKey) for which a new kernel has been registered
# The keys in the set are of the form `namespace + "/" + op_name + "/" + dispatch_key`.
# This set is maintained to ensure that two libraries don't try to override the exact same functionality to avoid
# libraries calling into kernels not intended to be called.
_impls: Set[str] = set()

# prim is reserved by TorchScript interpreter
_reserved_namespaces = ['prim']

[docs]def fallthrough_kernel(): """ A dummy function to pass to ``Library.impl`` in order to register a fallthrough. """ raise NotImplementedError("fallthrough_kernel() should never be called.")
[docs]class Library: """ A class to create libraries that can be used to register new operators or override operators in existing libraries from Python. A user can optionally pass in a dispatch keyname if they only want to register kernels corresponding to only one specific dispatch key. To create a library to override operators in an existing library (with name ns), set the kind to "IMPL". To create a new library (with name ns) to register new operators, set the kind to "DEF". To create a fragment of a possibly existing library to register operators (and bypass the limitation that there is only one library for a given namespace), set the kind to "FRAGMENT". Args: ns: library name kind: "DEF", "IMPL" (default: "IMPL"), "FRAGMENT" dispatch_key: PyTorch dispatch key (default: "") """ def __init__(self, ns, kind, dispatch_key=""): if kind not in ('IMPL', 'DEF', 'FRAGMENT'): raise ValueError("Unsupported kind: ", kind) if ns in _reserved_namespaces and (kind == "DEF" or kind == 'FRAGMENT'): raise ValueError(ns, " is a reserved namespace. Please try creating a library with another name.") frame = traceback.extract_stack(limit=3)[0] filename, lineno = frame.filename, frame.lineno self.m = torch._C._dispatch_library(kind, ns, dispatch_key, filename, lineno) self.ns = ns self._op_impls: Set[str] = set() self.kind = kind self.dispatch_key = dispatch_key # Use a finalizer to setup the "destructor" instead of __del__. # Python __del__ can lead to weird things (globals and locals may already # be gone when __del__ actually gets called!). finalizers help the # situation because it lets us capture references and keeps them alive weakref.finalize(self, _del_library, _impls, self._op_impls) def __repr__(self): return f"Library(kind={self.kind}, ns={self.ns}, dispatch_key={self.dispatch_key})>"
[docs] def define(self, schema, alias_analysis=""): r'''Defines a new operator and its semantics in the ns namespace. Args: schema: function schema to define a new operator. alias_analysis (optional): Indicates if the aliasing properties of the operator arguments can be inferred from the schema (default behavior) or not ("CONSERVATIVE"). Returns: name of the operator as inferred from the schema. Example:: >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LIBRARY) >>> my_lib = Library("foo", "DEF") >>> my_lib.define("sum(Tensor self) -> Tensor") ''' # This is added because we also want to disallow PURE_FUNCTION alias analysis which is a valid # AliasAnalysis type in C++ if alias_analysis not in ["", "FROM_SCHEMA", "CONSERVATIVE"]: raise RuntimeError(f"Invalid alias_analysis type {alias_analysis}") return self.m.define(schema, alias_analysis)
[docs] def impl(self, op_name, fn, dispatch_key=''): r'''Registers the function implementation for an operator defined in the library. Args: op_name: operator name (along with the overload) or OpOverload object. fn: function that's the operator implementation for the input dispatch key or :func:`~fallthrough_kernel` to register a fallthrough. dispatch_key: dispatch key that the input function should be registered for. By default, it uses the dispatch key that the library was created with. Example:: >>> my_lib = Library("aten", "IMPL") >>> def div_cpu(self, other): >>> return self * (1 / other) >>> my_lib.impl("div.Tensor", div_cpu, "CPU") ''' if not callable(fn): raise TypeError(f"Input function is required to be a callable but found type {type(fn)}") if dispatch_key == '': dispatch_key = self.dispatch_key if isinstance(op_name, str): name = op_name elif isinstance(op_name, OpOverload): name = overload_name = op_name._schema.overload_name if overload_name != '': name = name + '.' + overload_name else: raise RuntimeError("impl should be passed either a name or an OpOverload object as the first argument") key = self.ns + "/" + name.split("::")[-1] + "/" + dispatch_key if key in _impls: # TODO: in future, add more info about where the existing function is registered (this info is # today already returned by the C++ warning when impl is called but we error out before that) raise RuntimeError("This is not allowed since there's already a kernel registered from python overriding {}" "'s behavior for {} dispatch key and {} namespace.". format(name.split("::")[-1], dispatch_key, self.ns)) if dispatch_key == "Meta": dispatcher_op_name = name if '::' not in dispatcher_op_name: dispatcher_op_name = f'{self.ns}::{dispatcher_op_name}' # Internally, we shouldn't be registering meta kernels for any operators that # have CompositeImplicitAutograd kernels. # Instead, we should be letting those decompositions run, and writing meta kernels # only for the base operators. if torch._C._dispatch_has_kernel_for_dispatch_key(dispatcher_op_name, "CompositeImplicitAutograd"): raise RuntimeError( f"We should not register a meta kernel directly to the operator '{name}'," " because it has a CompositeImplicitAutograd kernel in core." " Instead we should let the operator decompose, and ensure that we have meta kernels" " for the base ops that it decomposes into.") self.m.impl(name, dispatch_key if dispatch_key != "" else "CompositeImplicitAutograd", fn) _impls.add(key) self._op_impls.add(key)
def _del_library(captured_impls, op_impls): captured_impls -= op_impls # decorator to register python functions for library ops # Note: this decorator API should remain consistent with `Library.impl` API def impl(lib, name, dispatch_key=""): def wrap(f): lib.impl(name, f, dispatch_key) return f return wrap def define(lib, schema, alias_analysis=""): def wrap(f): name = lib.define(schema, alias_analysis) lib.impl(name, f) return f return wrap


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources