torch.library

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

1. 创建新库

   • 允许您通过指定适当的分派键为各种后端和功能注册新的运算符和内核。例如

     • 考虑在新创建的 namespace 中注册新运算符 。您可以使用 API 访问此运算符，并通过调用 .您还可以通过调用 来访问特定的已注册重载。addfootorch.opstorch.ops.foo.addtorch.ops.foo.add.{overload_name}

     • 如果您为此运算符的调度键注册了新内核，则将为 CUDA 张量输入调用您的自定义定义函数。CUDA

   • 这可以通过创建 kind 的 Library 类对象来完成。"DEF"

2. 扩展现有的 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”。 :p aram ns：库名称 :p aram kind： “DEF”， “IMPL” （默认： “IMPL”） :p aram dispatch_key：PyTorch 调度密钥（默认值：“”）

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

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

参数

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

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

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

例：：

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

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

• torch.library.impl()

• torch.library.define()