torch.library¶

Python 操作员注册 API 提供了扩展 PyTorch 核心库的功能的运算符与用户定义的运算符。目前，这可以通过两种方式完成：

创建新库
- 允许您通过指定适当的分派键为各种后端和功能注册新的运算符和内核。例如
  - 考虑在新创建的 namespace 中注册新运算符。您可以使用 API 访问此运算符，并通过调用 .您还可以通过调用来访问特定的已注册重载。addfootorch.opstorch.ops.foo.addtorch.ops.foo.add.{overload_name}
  - 如果您为此运算符的调度键注册了新内核，则将为 CUDA 张量输入调用您的自定义定义函数。CUDA
- 这可以通过创建 kind 的 Library 类对象来完成。"DEF"
扩展现有的 C++ 库（例如 aten）
- 允许您通过指定适当的分派键，为与各种后端和功能对应的现有运算符注册内核。
- 这可能有助于填补对通过 dispatch key 实现的功能的参差不齐的 operator 支持。例如。
  - 您可以添加对 Meta Tensor 的运算符支持（通过将 function 注册到 dispatch key）。Meta
- 这可以通过创建 kind 的 Library 类对象来完成。"IMPL"

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

警告

Dispatcher 是一个复杂的 PyTorch 概念，对 Dispatcher 有充分的了解至关重要以便能够使用此 API 执行任何高级操作。此博客文章是了解 Dispatcher 的良好起点。

类 torch.library 中。库（ns， kind， dispatch_key=''）[来源]¶

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

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

参数

define（schema， alias_analysis=''）[来源]¶

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

参数

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

返回

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

例：：

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

impl（op_name， fn， dispatch_key=''）[来源]¶

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

参数

例：：

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

我们还添加了一些函数装饰器，以便于为运算符注册函数：

文档