(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](https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html)
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} More details on mutation, aliasing, and transforms
`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](https://colab.research.google.com/drive/1RRhSfk7So3Cn02itzLWE9K4Fam-8U011?usp=sharing).

```{eval-rst}
.. autoclass:: torch.library.Library
  :members:

.. autofunction:: fallthrough_kernel

.. autofunction:: define

.. autofunction:: impl
```