Table of Contents
Shortcuts

torch.library

torch.library 是一组用于扩展 PyTorch 核心库操作符的 API 集合。它包含创建新自定义操作符的实用工具，以及扩展使用 PyTorch 的 C++ 操作符注册 API 定义的操作符（例如 aten 操作符）。

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

使用 torch.library.define() 来定义新的自定义操作符。使用实现方法，例如 torch.library.impl() 和 func:torch.library.impl_abstract，来为任何操作符添加实现（它们可能已经使用 torch.library.define() 创建，或者通过 PyTorch 的 C++ 操作符注册 API 创建）。

torch.library.define(qualname, schema, *, lib=None, tags=())[源代码]
torch.library.define(lib, schema, alias_analysis='')

定义一个新运算符。

在 PyTorch 中，定义一个操作（简称“操作符”）是一个两步过程： - 我们需要定义操作（通过提供操作名称和模式） - 我们需要实现操作如何与各种 PyTorch 子系统（如 CPU/CUDA 张量、自动求导等）交互的行为。

此入口点定义了自定义操作符（第一步）， 然后您必须通过调用各种 impl_* API 来执行第二步，例如 torch.library.impl() 或 torch.library.impl_abstract()。

Parameters

  • qualname (str) – 操作符的限定名称。应该是一个类似于“namespace::name”的字符串，例如“aten::sin”。 PyTorch中的操作符需要一个命名空间来避免名称冲突；一个给定的操作符只能创建一次。 如果你正在编写一个Python库，我们建议命名空间为你顶级模块的名称。

  • schema (str) – 操作符的schema。例如，“(Tensor x) -> Tensor”表示一个接受一个Tensor并返回一个Tensor的操作符。它不包含操作符名称（该名称在qualname中传递）。

  • lib (可选[Library]) – 如果提供，该操作符的生命周期将与Library对象的生命周期绑定。

  • 标签 (标签 | 序列[标签]) – 一个或多个 torch.Tag 应用于该操作符。为操作符添加标签会改变操作符在各种 PyTorch 子系统下的行为；请在使用前仔细阅读 torch.Tag 的文档。

Example::
>>> import torch
>>> import numpy as np
>>>
>>> # 定义操作符
>>> torch.library.define("mylib::sin", "(Tensor x) -> Tensor")
>>>
>>> # 为操作符添加实现
>>> @torch.library.impl("mylibrary::sin", "cpu")
>>> def f(x):
>>>     return torch.from_numpy(np.sin(x.numpy()))
>>>
>>> # 从torch.ops调用新操作符
>>> x = torch.randn(3)
>>> y = torch.ops.mylib.sin(x)
>>> assert torch.allclose(y, x)
torch.library.impl(qualname, types, func=None, *, lib=None)[源代码]
torch.library.impl(lib, name, dispatch_key='')

为该操作员注册一个设备类型的实现。

您可以将“default”传递给 types，以将此实现注册为所有设备类型的默认实现。 请仅在实现真正支持所有设备类型时使用此选项； 例如，如果它是由内置的 PyTorch 运算符组成的，则这是成立的。

一些有效的类型包括：“cpu”、“cuda”、“xla”、“mps”、“ipu”、“xpu”。

Parameters

  • qualname (str) – 应该是一个看起来像“namespace::operator_name”的字符串。

  • 类型 (str | 序列[str]) – 要注册实现的目标设备类型。

  • lib (可选[Library]) – 如果提供，此注册的生命周期将与Library对象的生命周期绑定。

示例

>>> import torch
>>> import numpy as np
>>>
>>> # 定义操作符
>>> torch.library.define("mylibrary::sin", "(Tensor x) -> Tensor")
>>>
>>> # 为cpu设备添加实现
>>> @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)[源代码]

为此操作符注册一个抽象实现。

“抽象实现”指定了此运算符在携带无数据的张量上的行为。给定一些具有特定属性（大小/步幅/存储偏移/设备）的输入张量，它指定了输出张量的属性。

抽象实现具有与操作符相同的签名。它同时适用于FakeTensors和元张量。要编写抽象实现，假设操作符的所有张量输入都是常规的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
>>>
>>> # 示例1：一个没有数据依赖输出形状的操作符
>>> 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
>>>
>>> # 示例2：一个具有数据依赖输出形状的操作符
>>> torch.library.define("mylib::custom_nonzero", "(Tensor x) -> Tensor")
>>>
>>> @torch.library.impl_abstract("mylib::custom_nonzero")
>>> def custom_nonzero_abstract(x):
>>>     # 非零元素的数量是数据依赖的。
>>>     # 由于我们不能在抽象实现中查看数据，
>>>     # 我们使用ctx对象来构造一个新的symint，
>>>     # 表示数据依赖的大小。
>>>     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 对象。

调用 get_ctx() 仅在抽象实现内部有效 （参见 torch.library.impl_abstract() 以获取更多使用细节。

Return type

抽象实现上下文

低级 API

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

警告

低级操作符注册API和PyTorch调度器是PyTorch中一个复杂的概念。我们建议您尽可能使用更高级别的API（不需要torch.library.Library对象）。 这篇博客文章 <http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/>`_ 是了解PyTorch调度器的一个很好的起点。

一个教程，指导您通过一些示例了解如何使用此API，可在Google Colab上获取。

class torch.library.Library(ns, kind, dispatch_key='')[源代码]

一个用于创建库的类，该库可以从Python中注册新操作符或覆盖现有库中的操作符。 用户可以选择传递一个调度键名，如果他们只想注册与特定调度键对应的核函数。

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

Parameters

  • ns – 库名称

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

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

define(schema, alias_analysis='', *, tags=())[源代码]

在ns命名空间中定义一个新操作符及其语义。

Parameters

  • schema – 用于定义新操作符的函数架构。

  • alias_analysis (可选) – 指示操作符参数的别名属性是否可以从模式中推断出来（默认行为）或不推断（“CONSERVATIVE”）。

  • 标签 (标签 | 序列[标签]) – 一个或多个 torch.Tag 应用于该操作符。为操作符添加标签会改变操作符在各种 PyTorch 子系统下的行为；请在使用前仔细阅读 torch.Tag 的文档。

Returns

从模式推断出的操作员名称。

Example::
>>> my_lib = Library("foo", "DEF")
>>> my_lib.define("sum(Tensor self) -> Tensor")
impl(op_name, fn, dispatch_key='')[源代码]

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

Parameters

  • op_name – 操作符名称（连同重载）或 OpOverload 对象。

  • fn – 用于输入调度键的运算符实现的函数，或注册回退的fallthrough_kernel()。

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

Example::
>>> 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()[源代码]

一个传递给 Library.impl 的虚拟函数，用于注册一个回退处理。

优云智算