torch.library

PyTorch 操作符注册 API 提供了将用户自定义操作符扩展到 PyTorch 核心操作符库中的能力。目前，这可以通过两种方式实现：

1. 创建新的库

   • 允许您通过指定适当的分发密钥，为各种后端和功能注册新操作符和内核。例如,

     • Consider registering a new operator add in your newly created namespace foo. You can access this operator using the torch.ops API and calling into by calling torch.ops.foo.add. You can also access specific registered overloads by calling torch.ops.foo.add.{overload_name}.

     • If you registered a new kernel for the CUDA dispatch key for this operator, then your custom defined function will be called for CUDA tensor inputs.

   • 这可以通过创建"DEF"种类型的库类对象来实现。

2. 扩展现有的 C++ 库（例如 aten）

   • 允许你通过指定适当的调度键，为各种后端和功能对应的现有操作符注册内核。

   • 这在通过调度键实现的功能需要填补操作符支持不完整时可能会派上用场。例如，

     • You can add operator support for Meta Tensors (by registering function to the Meta dispatch key).

   • 这可以通过创建"IMPL"种类型的库类对象来实现。

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

警告

调度器是PyTorch中的一个复杂概念，对调度器有深入的理解对于能够使用此API进行高级操作至关重要。 这篇博客文章 是一个学习调度器的好起点。

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

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

Parameters

• ns – 库名称

• 类型 – “DEF”, “IMPL” (默认: “IMPL”)

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

define(schema, alias_analysis='')

根据一个模式定义一个新的操作符。 此外，还可以选择性地传入 alias_analysis 参数，以指示参数的别名属性是否可以从模式中推断出来（默认行为）或不推断（“CONSERVATIVE”）。

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

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

• torch.library.impl()

• torch.library.define()