Spaces:
Sleeping
Sleeping
| import inspect | |
| from torch._custom_op.impl import ( | |
| _custom_op_with_schema, | |
| _find_custom_op, | |
| infer_schema, | |
| parse_qualname, | |
| validate_namespace, | |
| ) | |
| from torch.library import get_ctx | |
| __all__ = [ | |
| "custom_op", | |
| "impl", | |
| "impl_abstract", | |
| "get_ctx", | |
| "impl_save_for_backward", | |
| "impl_backward", | |
| ] | |
| def custom_op(qualname, func_or_schema=None): | |
| r"""Register a new custom operator | |
| In PyTorch, defining an op (short for "operator") is a two step-process: | |
| - we need to define the op (by providing an operator name and schema) | |
| - we need to implement behavior for how the operator interacts with | |
| various PyTorch subsystems, like CPU/CUDA Tensors, Autograd, etc. | |
| This entrypoint defines the custom operator (the first step) | |
| you must then perform the second step by calling various | |
| ``impl_*`` APIs. | |
| This API may be used as a decorator (see examples). | |
| For a detailed guide on custom ops, please see | |
| https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk | |
| Arguments: | |
| qualname (str): Should be a string that looks like | |
| "namespace::operator_name". Operators in PyTorch need a namespace to | |
| avoid name collisions; a given operator may only be created once. | |
| If you are writing a Python library, we recommend the namespace to | |
| be the name of your top-level module. | |
| func_or_schema (Union[Callable, str]): Each PyTorch operator needs a | |
| schema that tells PyTorch the types of the inputs/outputs. | |
| If this is a Callable, we will automatically infer the schema from | |
| the type annotations on the function (see examples). Otherwise, | |
| if you don't want to use type annotations, you may provide us the | |
| schema string. | |
| Example:: | |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) | |
| >>> import torch | |
| >>> import numpy as np | |
| >>> from torch import Tensor | |
| >>> | |
| >>> # Step 1: define the custom op. | |
| >>> # We need to provide the API a "prototype function" | |
| >>> # (a function that returns NotImplementedError), from which | |
| >>> # we will infer the types of the inputs and outputs. | |
| >>> @torch._custom_ops.custom_op("mylibrary::numpy_sin") | |
| >>> def numpy_sin(x: Tensor) -> Tensor: | |
| >>> raise NotImplementedError() | |
| >>> | |
| >>> # The custom op is now accessible via the torch.ops module: | |
| >>> torch.ops.mylibrary.numpy_sin | |
| >>> | |
| >>> # Step 2: Register an implementation for various PyTorch subsystems | |
| >>> | |
| >>> # Register an implementation for CPU tensors | |
| >>> @torch._custom_ops.impl("mylibrary::numpy_sin", device_types="cpu") | |
| >>> def numpy_sin_impl_cpu(x): | |
| >>> return torch.from_numpy(np.sin(x.numpy())) | |
| >>> | |
| >>> # Register an implementation for CUDA tensors | |
| >>> @torch._custom_ops.impl("mylibrary::numpy_sin", device_types="cuda") | |
| >>> def numpy_sin_impl_cuda(x): | |
| >>> return torch.from_numpy(np.sin(x.cpu().numpy())).to(x.device) | |
| >>> | |
| >>> x = torch.randn(3) | |
| >>> torch.ops.mylibrary.numpy_sin(x) # calls numpy_sin_impl_cpu | |
| >>> | |
| >>> x_cuda = x.cuda() | |
| >>> torch.ops.mylibrary.numpy_sin(x) # calls numpy_sin_impl_cuda | |
| """ | |
| ns, name = parse_qualname(qualname) | |
| validate_namespace(ns) | |
| def inner(func): | |
| if not inspect.isfunction(func): | |
| raise ValueError( | |
| f"custom_op(...)(func): Expected `func` to be a Python " | |
| f"function, got: {type(func)}" | |
| ) | |
| if func.__name__ != name: | |
| raise ValueError( | |
| f"custom_op(qualname='{qualname}', ...)(func): expected `func` " | |
| f"to have name '{name}' but got '{func.__name__}'. " | |
| f"Please either change the name of `func` or the qualname that " | |
| f"is passed to `custom_op`" | |
| ) | |
| schema = infer_schema(func) | |
| _custom_op_with_schema(qualname, schema) | |
| return func | |
| if func_or_schema is None: | |
| return inner | |
| if isinstance(func_or_schema, str): | |
| _custom_op_with_schema(qualname, func_or_schema) | |
| else: | |
| return inner(func_or_schema) | |
| def impl(qualname, *, device_types=("cpu", "cuda"), func=None): | |
| r"""Register an implementation for a device type for this custom op. | |
| If the op is passed multiple Tensor inputs with different device | |
| types, it will dispatch to the registered implementation for the highest | |
| priority device type among those present. | |
| The supported device types, in order of priority, are {'cuda', 'cpu'}. | |
| This API may be used as a decorator (see examples). | |
| For a detailed guide on custom ops, please see | |
| https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk | |
| Arguments: | |
| device_types (str or Iterable[str]): the device type(s) to register the function for. | |
| Example:: | |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) | |
| >>> import torch | |
| >>> import numpy as np | |
| >>> from torch import Tensor | |
| >>> | |
| >>> # Step 1: define the custom op. | |
| >>> # We need to provide the API a "prototype function" | |
| >>> # (a function that returns NotImplementedError), from which | |
| >>> # we will infer the types of the inputs and outputs. | |
| >>> @torch._custom_ops.custom_op("mylibrary::numpy_cos") | |
| >>> def numpy_cos(x: Tensor) -> Tensor: | |
| >>> raise NotImplementedError() | |
| >>> | |
| >>> # The custom op is now accessible via the torch.ops module: | |
| >>> torch.ops.mylibrary.numpy_cos | |
| >>> | |
| >>> # Step 2: Register an implementation for various PyTorch subsystems | |
| >>> | |
| >>> # Register an implementation for CPU tensors | |
| >>> @torch._custom_ops.impl("mylibrary::numpy_cos", device_types="cpu") | |
| >>> def numpy_cos_impl_cpu(x): | |
| >>> return torch.from_numpy(np.cos(x.numpy())) | |
| >>> | |
| >>> # Register an implementation for CUDA tensors | |
| >>> @torch._custom_ops.impl("mylibrary::numpy_cos", device_types="cuda") | |
| >>> def numpy_cos_impl_cuda(x): | |
| >>> return torch.from_numpy(np.cos(x.cpu().numpy())).to(x.device) | |
| >>> | |
| >>> x = torch.randn(3) | |
| >>> torch.ops.mylibrary.numpy_cos(x) # calls numpy_cos_impl_cpu | |
| >>> | |
| >>> x_cuda = x.cuda() | |
| >>> torch.ops.mylibrary.numpy_cos(x) # calls numpy_cos_impl_cuda | |
| """ | |
| def inner(func): | |
| custom_op = _find_custom_op(qualname, also_check_torch_library=True) | |
| custom_op.impl(device_types, _stacklevel=3)(func) | |
| return func | |
| if func is None: | |
| return inner | |
| return inner(func) | |
| def impl_abstract(qualname, *, func=None): | |
| r"""Register an abstract implementation for this operator. | |
| An "abstract implementation" specifies the behavior of this operator on | |
| Tensors that carry no data. Given some input Tensors with certain properties | |
| (sizes/strides/storage_offset/device), it specifies what the properties of | |
| the output Tensors are. | |
| The abstract implementation has the same signature as the operator. | |
| It is run for both FakeTensors and meta tensors. To write an abstract | |
| implementation, assume that all Tensor inputs to the operator are | |
| regular CPU/CUDA/Meta tensors, but they do not have storage, and | |
| you are trying to return regular CPU/CUDA/Meta tensor(s) as output. | |
| The abstract implementation must consist of only PyTorch operations | |
| (and may not directly access the storage or data of any input or | |
| intermediate Tensors). | |
| This API may be used as a decorator (see examples). | |
| For a detailed guide on custom ops, please see | |
| https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk | |
| Examples:: | |
| >>> import numpy as np | |
| >>> from torch import Tensor | |
| >>> | |
| >>> # Example 1: an operator without data-dependent output shape | |
| >>> @torch._custom_ops.custom_op("mylibrary::custom_linear") | |
| >>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor: | |
| >>> raise NotImplementedError() | |
| >>> | |
| >>> @torch._custom_ops.impl_abstract("mylibrary::custom_linear") | |
| >>> def custom_linear_abstract(x, weight): | |
| >>> assert x.dim() == 2 | |
| >>> assert weight.dim() == 2 | |
| >>> assert bias.dim() == 1 | |
| >>> assert x.shape[1] == weight.shape[1] | |
| >>> assert weight.shape[0] == bias.shape[0] | |
| >>> assert x.device == weight.device | |
| >>> | |
| >>> return (x @ weight.t()) + bias | |
| >>> | |
| >>> # Example 2: an operator with data-dependent output shape | |
| >>> @torch._custom_ops.custom_op('mylibrary::custom_nonzero') | |
| >>> def custom_nonzero(x: Tensor) -> Tensor: | |
| >>> ... | |
| >>> | |
| >>> @torch._custom_ops.impl_abstract("mylibrary::custom_nonzero") | |
| >>> def custom_nonzero_abstract(x): | |
| >>> # Number of nonzero-elements is data-dependent. | |
| >>> # Since we cannot peek at the data in an abstract impl, | |
| >>> # we use the ctx object to construct a new symint that | |
| >>> # represents the data-dependent size. | |
| >>> ctx = torch._custom_ops.get_ctx() | |
| >>> nnz = ctx.create_unbacked_symint() | |
| >>> shape = [x.dim(), nnz] | |
| >>> result = x.new_empty(shape, dtype=torch.long) | |
| >>> return result | |
| >>> | |
| >>> @torch._custom_ops.impl("mylibrary::custom_nonzero") | |
| >>> def custom_nonzero_impl(x): | |
| >>> x_np = to_numpy(x) | |
| >>> res = np.stack(np.nonzero(x_np), axis=1) | |
| >>> # unbacked symbolic ints in PyTorch must be >= 2, so we | |
| >>> # constrain the range to at least 2 | |
| >>> if res.shape[0] <= 1: | |
| >>> raise RuntimeError("not supported") | |
| >>> return torch.tensor(res, device=x.device) | |
| """ | |
| import torch.library | |
| return torch.library.impl_abstract(qualname, func, _stacklevel=2) | |
| def impl_save_for_backward(qualname, *, func=None): | |
| r"""Register a function that tells us what to save for backward. | |
| Please see :func:`impl_backward` for more details. | |
| """ | |
| def inner(func): | |
| custom_op = _find_custom_op(qualname, also_check_torch_library=True) | |
| custom_op.impl_save_for_backward(_stacklevel=3)(func) | |
| return func | |
| if func is None: | |
| return inner | |
| return inner(func) | |
| def impl_backward(qualname, output_differentiability=None, *, func=None): | |
| r"""Registers a backward formula for an operator. | |
| In order for an operator to work with autograd, you need to register | |
| a backward formula. There are two pieces to this: | |
| 1. You must give us a function to specify what to save for backward. | |
| Call this the "save for backward" function. | |
| 2. You must give us a function that computes gradients. Call this the | |
| "backward" function. | |
| Use `impl_save_for_backward` to define a "save for backward" function | |
| that specifies what gets saved for backward. The function should accept | |
| two arguments ``(inputs, output)`` and return the quantities to be saved | |
| for backward. | |
| During runtime, when you call the operator in a forwards pass, PyTorch | |
| will invoke the "save for backward" function with the inputs and output | |
| of the operator. | |
| Use `impl_backward` to define the "backward" function. The backward | |
| function must accept ``(ctx, saved, *grads)``: | |
| - ``ctx`` is a context object where we may provide information | |
| - ``saved`` is exactly what gets returned from the "save for backward" | |
| function | |
| - ``grads`` is one or more gradients. The number of gradients matches | |
| the number of outputs of the operator. | |
| The backward function must return a dict that maps the name of | |
| an input to the operator to its corresponding gradient. All inputs that | |
| were declared to be Tensors in the operator definition must be accounted | |
| for in the dict. The gradient may be a Tensor or None. | |
| For a detailed guide on custom ops, please see | |
| https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk | |
| """ | |
| def inner(func): | |
| custom_op = _find_custom_op(qualname, also_check_torch_library=True) | |
| custom_op.impl_backward(output_differentiability, _stacklevel=3)(func) | |
| return func | |
| if func is None: | |
| return inner | |
| return inner(func) | |
| def _destroy(qualname): | |
| """De-registers a custom op. For testing purposes only""" | |
| custom_op = _find_custom_op(qualname) | |
| custom_op._destroy() | |