torch.library

torch.library 是用于扩展 PyTorch 核心库的 API 集合 的运算符。它包含用于创建新自定义运算符的实用程序 以及扩展使用 PyTorch 的 C++ 运算符定义的运算符 注册 API（例如 aten 运算符）。

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

用于定义新的自定义运算符。使用 impl 方法，例如 和 func：torch.library.impl_abstract 添加实现 对于任何运算符（它们可能是使用 或 通过 PyTorch 的 C++ 操作员注册 API）。

torch.library 中。define（qualname， schema， *， lib=None， tags=（））

torch.library 中。define（lib， 架构， alias_analysis='')

定义 new 运算符。

在 PyTorch 中，定义 op（“operator”的缩写）是一个两步过程： - 我们需要定义 OP（通过提供 Operator Name 和 schema） - 我们需要实现 Operator 如何与 各种 PyTorch 子系统，如 CPU/CUDA 张量、Autograd 等。

此入口点定义自定义运算符（第一步） 然后，您必须通过调用各种 API 来执行第二步，例如 或 。impl_*

参数

• qualname （str） – 运算符的限定名称。应该是 一个看起来像 “namespace：：name” 的字符串，例如 “aten：：sin”。 PyTorch 中的 Operator 需要一个命名空间来 避免名称冲突;给定的运算符只能创建一次。 如果您正在编写 Python 库，我们建议将命名空间设置为 是顶级模块的名称。

• schema （str） – 运算符的架构。例如“（Tensor x） -> Tensor” 对于接受一个 Tensor 并返回一个 Tensor 的运算。它确实 不包含运算符名称（在 中传递）。qualname

• lib （Optional[Library]） – 如果提供，则此运算符的生命周期 将与 Library 对象的生命周期相关联。

• 标签 （Tag | Sequence[Tag]） – 一个或多个火把。应用于此的标签 算子。标记运算符会更改运算符的行为 在各种 PyTorch 子系统下;请阅读 torch。应用前请仔细标记。

例：：

>>> 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=无， *， lib=无）

torch.library 中。impl（lib， name， dispatch_key='')

为此 Operator 的设备类型注册一个 implementation 。

你可以传递 “default” for 将此实现注册为 default 实现。 请仅在实现真正支持所有设备类型时才使用此项; 例如，如果它是内置 PyTorch 运算符的组合，则为 true。types

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

参数

• qualname （str） – 应为类似于 “namespace：：operator_name” 的字符串。

• 类型 （str | Sequence[str]） – 要将 impl 注册到的设备类型。

• lib （Optional[Library]） – 如果提供，则此注册的生命周期 将与 Library 对象的生命周期相关联。

例子

>>> 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）

为此运算符注册一个 abstract implementation 。

“抽象实现” 指定此运算符在 不携带数据的 Tensor。给定一些具有特定属性的输入 Tensor （sizes/strides/storage_offset/device） 中，它指定了 output Tensors 是。

抽象实现具有与 operator 相同的签名。 它同时针对 FakeTensor 和 meta Tensors 运行。编写摘要 implementation 中，假设算子的所有 Tensor input 都是 常规 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 对象。

调用仅在抽象 impl 内部有效 （有关更多使用详情，请参阅。get_ctx()

返回类型

抽象 ImplCtx

低级 API

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

警告

低级操作员注册 API 和 PyTorch Dispatcher 是一个 复杂的 PyTorch 概念。我们建议您使用上述更高级别的 API （不需要 torch.library.Library 对象）。 这篇博文 <http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/>'_ 是了解 PyTorch Dispatcher 的良好起点。

Google Colab 上提供了指导您完成有关如何使用此 API 的一些示例的教程。

类 torch.library 中。库（ns， kind， dispatch_key=''）

用于创建可用于注册新运算符或 override Python 中现有库中的运算符。 如果用户只想注册，则可以选择传入 dispatch keyname kernels 只对应于一个特定的 dispatch key。

要创建一个库来覆盖现有库（名称为 ns）中的运算符，请将 kind 设置为 “IMPL”。 要创建一个新库（名称为 ns）来注册新的运算符，请将 kind 设置为 “DEF”。 要创建可能存在的库的片段以注册运算符（并绕过 给定命名空间只有一个库的限制），将 kind 设置为 “片段”。

参数

• NS – 库名称

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

• dispatch_key – PyTorch 调度密钥（默认值：“”）

define（schema， alias_analysis=''， *， tags=（）））

在 ns 命名空间中定义 new 运算符及其语义。

参数

• schema – 用于定义新运算符的函数 schema。

• alias_analysis （可选） – 指示运算符参数的别名属性是否可以 从架构 （默认行为） 推断出 （“CONSERVATIVE”） 或非 （“CONSERVATIVE”） 。

• 标签 （Tag | Sequence[Tag]） – 一个或多个火把。应用于此的标签 算子。标记运算符会更改运算符的行为 在各种 PyTorch 子系统下;请阅读 torch。应用前请仔细标记。

返回

从架构推断的运算符的名称。

例：：

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

impl（op_name， fn， dispatch_key=''）

为库中定义的运算符注册函数实现。

参数

• op_name – 运算符名称（以及重载）或 OpOverload 对象。

• fn – 作为 Input Dispatch Key 或注册 fallthrough 的运算符实现的函数。

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

例：：

>>> 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（）

一个虚拟函数，用于注册 fallthrough。Library.impl