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 的良好起点。

class （ns， kind， dispatch_key=''）torch.library.Library

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

参数

• NS – 库名称

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

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

define(图式， alias_analysis=''）

采用 schema 来定义 new 运算符。 此外，还可以选择采用 alias_analysis 参数来指示参数的别名属性 可以从架构 （默认行为） 或非架构 （“CONSERVATIVE”） 推断。

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

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

• torch.library.impl()

• torch.library.define()