torch.library

torch.library 是一组用于扩展 PyTorch 核心运算符库的 API 集合。它包含用于创建新的自定义运算符的工具，以及扩展通过 PyTorch 的 C++ 运算符注册 API（例如 aten 运算符）定义的运算符的功能。

有关有效使用这些API的详细指南，请参阅 这份gdoc

使用 torch.library.define() 来定义新的自定义操作符。使用 impl 方法，例如 torch.library.impl() 和 func:torch.library.impl_abstract，来为任何操作符添加实现（它们可能已经通过 torch.library.define() 或 PyTorch 的 C++ 操作符注册 API 创建）。

torch.library.define(qualname, schema, *, lib=None, tags=())

torch.library.define(lib, schema, alias_analysis='')

定义一个新的算子。

在 PyTorch 中，定义一个算子（op 的全称为“operator”）是一个两步过程： - 我们需要定义该算子（提供算子名称和模式） - 我们需要实现该算子与各种 PyTorch 子系统的交互行为，比如 CPU/CUDA 张量、自动微分等。

此入口定义了自定义操作符（第一步）， 然后您必须通过调用各种 impl_* API 来执行第二步，例如 torch.library.impl() 或 torch.library.impl_abstract()。

Parameters

• qualname (str) – 操作符的限定名。应为类似于“命名空间::名称”的字符串，例如“aten::sin”。 PyTorch中的操作符需要一个命名空间以避免名称冲突；给定的操作符只能创建一次。 如果您正在编写Python库，我们建议命名空间使用顶级模块的名称。

• 模式 (字符串) – 算子的模式。例如，对于接受一个张量并返回一个张量的操作，其模式为 “(Tensor x) -> Tensor”。 这里不包含操作名称（操作名称通过 qualname 传递）。

• 库 (可选[库]) – 如果提供，此操作符的生命周期将与库对象的生命周期绑定。

• 标签 (标签 | 序列[标签]) – 一个或多个 torch.Tag，应用于此操作符。标记一个操作符会改变该操作符在各种 PyTorch 子系统下的行为；请在应用之前仔细阅读 torch.Tag 的文档。

Example::

>>> import torch
>>> import numpy as np
>>>
>>> # Define the operator
>>> torch.library.define("mylib::sin", "(Tensor x) -> Tensor")
>>>
>>> # Add implementations for the operator
>>> @torch.library.impl("mylibrary::sin", "cpu")
>>> def f(x):
>>>     return torch.from_numpy(np.sin(x.numpy()))
>>>
>>> # Call the new operator from torch.ops.
>>> x = torch.randn(3)
>>> y = torch.ops.mylib.sin(x)
>>> assert torch.allclose(y, x)

torch.library.impl(qualname, types, func=None, *, lib=None)

torch.library.impl(lib, name, dispatch_key='')

为该操作的设备类型注册一个实现。

您可以将“default”传递给types，以将此实现注册为所有设备类型的默认实现。 请仅在该实现真正支持所有设备类型时使用此选项；例如，如果它是内置的PyTorch操作符组合，则此条件成立。

一些有效的类型包括：“cpu”，“cuda”，“xla”，“mps”，“ipu”，“xpu”。

Parameters

• qualname (str) – 应该是一个类似“命名空间::操作符名称”的字符串。

• 类型 (str | Sequence[str]) – 要注册实现的设备类型。

• 库 (可选[库]) – 如果提供，此注册的生命周期将与库对象的生命周期相关联。

示例

>>> import torch
>>> import numpy as np
>>>
>>> # Define the operator
>>> torch.library.define("mylibrary::sin", "(Tensor x) -> Tensor")
>>>
>>> # Add implementations for the cpu device
>>> @torch.library.impl("mylibrary::sin", "cpu")
>>> def f(x):
>>>     return torch.from_numpy(np.sin(x.numpy()))
>>>
>>> x = torch.randn(3)
>>> y = torch.ops.mylibrary.sin(x)
>>> assert torch.allclose(y, x.sin())

torch.library.impl_abstract(qualname, func=None, *, lib=None, _stacklevel=1)

为该操作符注册一个抽象实现。

“抽象实现”指定了该操作符在不携带数据的张量上的行为。给定一些具有特定属性（尺寸/步幅/存储偏移/设备）的输入张量，它会指定输出张量的属性是什么。

抽象实现的签名与操作符相同。 它会同时对 FakeTensors 和 meta 张量运行。要编写一个抽象实现，请假设操作符的所有 Tensor 输入都是常规的 CPU/CUDA/Meta 张量，但它们没有存储空间，并且你试图返回常规的 CPU/CUDA/Meta 张量作为输出。 抽象实现必须仅由 PyTorch 操作组成（并且不能直接访问任何输入或中间张量的存储空间或数据）。

此 API 可用作装饰器（参见示例）。

如需详细了解自定义操作，请参阅 https://docs.google.com/document/d/1W–T6wz8IY8fOI0Vm8BF44PdBgs283QvpelJZWieQWQ/edit

示例

>>> import torch
>>> import numpy as np
>>> from torch import Tensor
>>>
>>> # Example 1: an operator without data-dependent output shape
>>> torch.library.define(
>>>     "mylib::custom_linear",
>>>     "(Tensor x, Tensor weight, Tensor bias) -> Tensor")
>>>
>>> @torch.library.impl_abstract("mylib::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.library.define("mylib::custom_nonzero", "(Tensor x) -> Tensor")
>>>
>>> @torch.library.impl_abstract("mylib::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.library.get_ctx()
>>>     nnz = ctx.new_dynamic_size()
>>>     shape = [nnz, x.dim()]
>>>     result = x.new_empty(shape, dtype=torch.int64)
>>>     return result
>>>
>>> @torch.library.impl("mylib::custom_nonzero", "cpu")
>>> def custom_nonzero_cpu(x):
>>>     x_np = x.numpy()
>>>     res = np.stack(np.nonzero(x_np), axis=1)
>>>     return torch.tensor(res, device=x.device)

torch.library.get_ctx()

get_ctx() 返回当前的 AbstractImplCtx 对象。

调用 get_ctx() 仅在抽象实现内部有效 （详见 torch.library.impl_abstract() 获取更多使用详情。

Return type

AbstractImplCtx

低级APIs

以下 API 是对 PyTorch 的 C++ 低级操作注册 API 的直接绑定。

警告

低级别的算子注册API和PyTorch调度器是PyTorch中的一个复杂概念。我们建议您在可能的情况下使用更高层次的API（这些API不需要torch.library.Library对象）。 这篇博客文章<http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/>是一个了解PyTorch调度器的好起点。

在Google Colab上提供了一个教程，通过一些示例向您展示如何使用此API。

class torch.library.Library(ns, kind, dispatch_key='')

一个类，用于创建可以从 Python 中使用的库，以注册新操作符或覆盖现有库中的操作符。 用户可以选择性地传入一个调度键名，以便仅注册与特定调度键相对应的内核。

要创建一个用于重载现有库（名称为 ns）操作符的库，请将类型设置为“IMPL”。 要创建一个新的库（名称为 ns）以注册新的操作符，请将类型设置为“DEF”。 要创建一个可能已存在的库片段以注册操作符（并绕过给定命名空间只能有一个库的限制），请将类型设置为“FRAGMENT”。

Parameters

• ns – 库名称

• 类型 – “DEF”，“IMPL”（默认：“IMPL”），“FRAGMENT”

• dispatch_key – PyTorch 调度键 (默认值：”“)

define(schema, alias_analysis='', *, tags=())

在命名空间 ns 中定义一个新的操作符及其语义。

Parameters

• 模式 – 定义新算子的功能模式。

• alias分析 (可选) – 表示操作数的别名属性是否可以从模式（默认行为）中推断出来，或者不能（“保守”）。

• 标签 (标签 | 序列[标签]) – 一个或多个 torch.Tag，应用于此操作符。标记一个操作符会改变该操作符在各种 PyTorch 子系统下的行为；请在应用之前仔细阅读 torch.Tag 的文档。

Returns

根据模式推断的操作符名称。

Example::

>>> my_lib = Library("foo", "DEF")
>>> my_lib.define("sum(Tensor self) -> Tensor")

impl(op_name, fn, dispatch_key='')

为库中定义的操作符注册函数实现。

Parameters

• op_name – 操作符名称（包括重载）或 OpOverload 对象。

• fn – 输入分发键的操作实现函数，或fallthrough_kernel() 以注册一个备用操作。

• dispatch_key – 输入函数应注册的分派键。默认情况下，它使用库创建时的分派键。

Example::

>>> my_lib = Library("aten", "IMPL")
>>> def div_cpu(self, other):
>>>     return self * (1 / other)
>>> my_lib.impl("div.Tensor", div_cpu, "CPU")

torch.library.fallthrough_kernel()

一个虚拟函数，用于传递给Library.impl以注册一个贯穿。