(torch-library-docs)=

torch.library

{eval-rst}

.. py:module:: torch.library
.. currentmodule:: torch.library

torch.library is a collection of APIs for extending PyTorch's core library of operators. It contains utilities for testing custom operators, creating new custom operators, and extending operators defined with PyTorch's C++ operator registration APIs (e.g. aten operators).

For a detailed guide on effectively using these APIs, please see PyTorch Custom Operators Landing Page for more details on how to effectively use these APIs.

Testing custom ops

Use {func}torch.library.opcheck to test custom ops for incorrect usage of the Python torch.library and/or C++ TORCH_LIBRARY APIs. Also, if your operator supports training, use {func}torch.autograd.gradcheck to test that the gradients are mathematically correct.

{eval-rst}

.. autofunction:: opcheck

Creating new custom ops in Python

Use {func}torch.library.custom_op to create new custom ops.

Choosing the kind of custom op

When defining a custom op, first decide what aliasing and mutation behavior the operator has. This choice determines the mutates_args, tags, return schema, and fake-kernel behavior.

text

Does the operator mutate any Tensor input?
  |
  +-- no --> Functional operator
  |         Use mutates_args=().
  |         Return new Tensor objects that do not alias inputs or each other.
  |
  +-- yes --> Does it mutate the first positional Tensor argument and return it?
        |
        +-- yes --> In-place operator
        |          Use tags=torch.Tag.inplace.
        |          Use mutates_args with exactly the first argument name.
        |          Return that same first argument.
        |
        +-- no --> Does it mutate keyword-only Tensor out= arguments and return them?
              |
              +-- yes --> out= operator
              |          Use tags=torch.Tag.out.
              |          Do not read from the out= tensors.
              |          Put mutable output tensors after * as keyword-only args.
              |          Return all out= tensors in declaration order.
              |
              +-- no --> Mutable operator
                         Use mutates_args for the mutated arguments.
                         Do not return mutated inputs or aliases of inputs.

Functional operators are the simplest and should be preferred when possible. They should not mutate inputs, and their outputs should be fresh values that do not alias any input or other output.

Use tags=torch.Tag.inplace only for the conventional in-place shape: the first positional argument is a Tensor, it is the only mutated argument, and the operator returns that same Tensor. For example, an inferred in-place custom op has a schema shaped like:

text

(Tensor(a!) x, ...) -> Tensor(a!)

Use tags=torch.Tag.out only for conventional out= variants: all mutable output tensors are keyword-only arguments, and the operator returns those outputs in declaration order. The implementation must write to out= tensors, but must not read from them. For example:

text

(Tensor x, *, Tensor(a!) out) -> Tensor(a!)
(Tensor x, *, Tensor(a!) out0, Tensor(b!) out1) -> (Tensor(a!), Tensor(b!))

In-place and out= custom ops get an autogenerated fake kernel from their tag and schema. Registering a fake kernel for these ops is usually unnecessary, though it is still allowed if you need to override the autogenerated behavior. out= custom ops follow the same autograd rule as built-in out= operators: they do not support automatic differentiation when any argument requires grad.

Arbitrary aliasing, such as returning an alias of an input that is not one of the in-place or out= patterns above, is not supported by custom_op. PyTorch compiler transforms such as functionalization need to reason precisely about mutations and aliases, so custom op aliasing is intentionally limited to the conventional forms described here. If your operator mutates inputs but is not an in-place or out= operator, model it as a mutable operator and return no aliases of inputs.

The same operation can have different contracts depending on how it handles its outputs:

python

from torch import Tensor


@torch.library.custom_op(
    "mylib::add_inplace",
    mutates_args={"x"},
    tags=torch.Tag.inplace,
)
def add_inplace(x: Tensor, y: Tensor) -> Tensor:
    x.add_(y)
    return x


@torch.library.custom_op(
    "mylib::add_out",
    mutates_args={"out"},
    tags=torch.Tag.out,
)
def add_out(x: Tensor, y: Tensor, *, out: Tensor) -> Tensor:
    out.copy_(x + y)
    return out


@torch.library.custom_op("mylib::add_mutate", mutates_args={"x"})
def add_mutate(x: Tensor, y: Tensor) -> None:
    x.add_(y)

{dropdown}

`custom_op` asks for a precise mutation and aliasing contract because PyTorch
uses that contract in FakeTensor, autograd, functionalization, and
`torch.compile`.

For a functional custom op, PyTorch assumes the operator does not mutate any
input and that returned tensors are fresh values. This is the easiest kind of op
for PyTorch transforms to reason about.

For an in-place custom op, `torch.Tag.inplace` gives PyTorch a stronger and more
specific contract: the first Tensor argument is mutated and the returned Tensor
is the same object. This lets PyTorch derive the fake behavior from the schema
instead of requiring a separate fake kernel.

For an `out=` custom op, `torch.Tag.out` gives PyTorch the corresponding
`out=`-variant contract: keyword-only output tensors are mutated and returned in
schema order. The `out=` tensors are output buffers only: the implementation
must not use their current values as inputs. Like built-in `out=` operators,
these custom ops do not support automatic differentiation when any argument
requires grad.

For other mutable custom ops, use `mutates_args` to name the mutated arguments,
but do not return mutated inputs or aliases of inputs. Arbitrary aliasing is
hard for transforms because functionalization rewrites mutating programs into
functional programs and must know exactly which returned tensors share storage
with which inputs. `custom_op` therefore supports the conventional in-place and
`out=` aliasing patterns directly, but does not model arbitrary view or alias
relationships.

{eval-rst}

.. autofunction:: custom_op
.. autofunction:: triton_op
.. autofunction:: wrap_triton

Extending custom ops (created from Python or C++)

Use the register.* methods, such as {func}torch.library.register_kernel and {func}torch.library.register_fake, to add implementations for any operators (they may have been created using {func}torch.library.custom_op or via PyTorch's C++ operator registration APIs).

{eval-rst}

.. autofunction:: register_kernel
.. autofunction:: register_autocast
.. autofunction:: register_autograd
.. autofunction:: register_fake
.. autofunction:: register_vmap
.. autofunction:: impl_abstract
.. autofunction:: get_ctx
.. autofunction:: register_torch_dispatch
.. autofunction:: infer_schema
.. autoclass:: torch._library.custom_ops.CustomOpDef
   :members: set_kernel_enabled
.. autofunction:: get_kernel

Low-level APIs

The following APIs are direct bindings to PyTorch's C++ low-level operator registration APIs.

{eval-rst}

.. warning:: The low-level operator registration APIs and the PyTorch Dispatcher are a complicated PyTorch concept. We recommend you use the higher level APIs above (that do not require a torch.library.Library object) when possible. `This blog post <http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/>`_ is a good starting point to learn about the PyTorch Dispatcher.

A tutorial that walks you through some examples on how to use this API is available on Google Colab.

{eval-rst}

.. autoclass:: torch.library.Library
  :members:

.. autofunction:: fallthrough_kernel

.. autofunction:: define

.. autofunction:: impl