torch.library

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

有关如何有效使用这些API的详细指南，请参见 PyTorch 自定义操作符登录页面 以了解更多关于如何有效使用这些API的信息。

测试自定义操作

使用 torch.library.opcheck() 测试自定义操作，以检查 Python torch.library 和/或 C++ TORCH_LIBRARY API 的错误用法。此外，如果您的操作支持训练，请使用 torch.autograd.gradcheck() 测试梯度的数学正确性。

torch.library.opcheck(op, args, kwargs=None, *, test_utils=('test_schema', 'test_autograd_registration', 'test_faketensor', 'test_aot_dispatch_dynamic'), raise_exception=True)

给定一个算子和一些样本参数，测试该算子是否正确注册。

也就是说，当你使用torch.library/TORCH_LIBRARY API创建自定义算子时，你需要指定该自定义算子的元数据（例如可变性信息），并且这些API要求你传递给它们的函数必须满足某些属性（例如在假/元/抽象内核中不访问数据指针）。opcheck测试这些元数据和属性。

具体来说，我们测试了以下内容：

• test_schema：如果模式与操作符的实现匹配。例如：如果模式指定张量会被修改，则我们检查实现是否确实修改了该张量。如果模式指定我们返回一个新的张量，则我们检查实现是否确实返回了一个新的张量（而不是现有的张量或现有张量的视图）。

• test_autograd_registration：如果操作符支持训练（autograd），我们会检查其 autograd 公式是否通过 torch.library.register_autograd 或者对一个或多个 DispatchKey::Autograd 键的手动注册进行了注册。任何基于其他 DispatchKey 的注册都可能导致未定义行为。

• test_faketensor：如果操作符具有 FakeTensor 内核（并且该内核是正确的）。FakeTensor 内核是操作符与 PyTorch 编译 API（torch.compile/export/FX）配合工作的必要条件（但非充分条件）。我们检查操作符是否注册了 FakeTensor 内核（有时也称为元内核），并且该内核是否正确。此测试会获取操作符在真实张量上运行的结果以及在 FakeTensor 上运行的结果，并检查它们的张量元数据（大小/步幅/数据类型/设备等）是否相同。

• test_aot_dispatch_dynamic：如果操作符在使用 PyTorch 编译 API（torch.compile/export/FX）时具有正确的行为。 此检查验证在急切模式 PyTorch 和 torch.compile 下的输出（以及适用的梯度）是否相同。 此测试是 test_faketensor 的超集，并且是一个端到端测试； 它还测试了操作符是否支持函数化，以及反向传递（如果存在）是否也支持 FakeTensor 和函数化。

为了获得最佳效果，请多次调用opcheck，并使用一组具有代表性的输入。如果您的操作符支持自动微分，请使用opcheck，并输入requires_grad = True；如果您的操作符支持多个设备（例如CPU和CUDA），请使用opcheck，并在所有支持的设备上进行输入。

Parameters

• 操作 (联合[OpOverload, OpOverloadPacket, 自定义操作定义]) – 操作符。必须是一个用 torch.library.custom_op() 装饰的函数，或者是 torch.ops.* 中的 OpOverload/OpOverloadPacket （例如 torch.ops.aten.sin, torch.ops.mylib.foo）

• 参数 (元组[任意类型, ...]) – 操作符的参数

• kwargs (可选[字典[字符串, 任意类型]]) – 传递给操作符的关键字参数

• test_utils (Union[str, Sequence[str]]) – 我们应该运行的测试。默认：全部运行。 示例：（“test_schema”，“test_faketensor”）

• raise_exception (bool) – 如果应在第一次出错时抛出异常。如果为False，我们将返回一个字典，其中包含每个测试是否通过的信息。

Return type

字典[字符串, 字符串]

警告

opcheck 和 torch.autograd.gradcheck() 测试不同的内容； opcheck 检查你对 torch.library API 的使用是否正确，而 torch.autograd.gradcheck() 检查你的自动微分公式是否数学上正确。请使用两者来测试支持梯度计算的自定义操作。

示例

>>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
>>> def numpy_mul(x: Tensor, y: float) -> Tensor:
>>>     x_np = x.numpy(force=True)
>>>     z_np = x_np * y
>>>     return torch.from_numpy(z_np).to(x.device)
>>>
>>> @numpy_mul.register_fake
>>> def _(x, y):
>>>     return torch.empty_like(x)
>>>
>>> def setup_context(ctx, inputs, output):
>>>     y, = inputs
>>>     ctx.y = y
>>>
>>> def backward(ctx, grad):
>>>     return grad * ctx.y, None
>>>
>>> numpy_mul.register_autograd(backward, setup_context=setup_context)
>>>
>>> sample_inputs = [
>>>     (torch.randn(3), 3.14),
>>>     (torch.randn(2, 3, device='cuda'), 2.718),
>>>     (torch.randn(1, 10, requires_grad=True), 1.234),
>>>     (torch.randn(64, 64, device='cuda', requires_grad=True), 90.18),
>>> ]
>>>
>>> for args in sample_inputs:
>>>     torch.library.opcheck(numpy_mul, args)

用Python创建新的自定义操作

使用 torch.library.custom_op() 创建新的自定义操作。

torch.library.custom_op(name, fn=None, /, *, mutates_args, device_types=None, schema=None)

将函数包装为自定义操作符。

您可能希望创建自定义运算的原因包括： - 将第三方库或自定义内核包装起来，以便与 PyTorch 子系统（如自动微分）一起使用。 - 防止 torch.compile/导出/FX 追踪窥探您的函数内部。

此 API 用作函数周围的装饰器（请参见示例）。 提供的函数必须具有类型提示；这些类型提示用于与 PyTorch 的各种子系统进行接口。

Parameters

• 名称 (str) – 自定义操作的名称，格式类似于“{命名空间}::{名称}”， 例如 “mylib::my_linear”。该名称在 PyTorch 子系统（如 torch.export、FX 图）中用作操作的稳定标识符。 为避免名称冲突，请使用您的项目名称作为命名空间； 例如，pytorch/fbgemm 中的所有自定义操作都使用 “fbgemm” 作为命名空间。

• mutates_args (可迭代对象[str] 或 "unknown") – 函数变异的参数名称。 这必须准确，否则行为将无法预测。如果为“unknown”，则悲观地假设操作符的所有输入都被变异。

• 设备类型 (None | str | Sequence[str]) – 该函数有效的设备类型。如果没有提供设备类型，则该函数将作为所有设备类型的默认实现。 示例：“cpu”，“cuda”。 当为不接受张量的操作符注册特定设备的实现时，我们要求该操作符必须有一个“device: torch.device 参数”。

• schema (无 | 字符串) – 操作符的模式字符串。如果为None （推荐），我们将从其类型注释中推断出操作符的模式。除非您有特定原因，否则我们建议让系统自动推断模式。 示例：“(Tensor x, int y) -> (Tensor, Tensor)”。

Return type

任何

注意

我们建议不要传递 schema 参数，而是让我们从类型注释中推断。 自己编写模式容易出错。如果您对我们对类型注释的解释不满意， 您可以提供自己的模式。有关如何编写模式字符串的更多信息，请参见 这里

Examples::

>>> import torch
>>> from torch import Tensor
>>> from torch.library import custom_op
>>> import numpy as np
>>>
>>> @custom_op("mylib::numpy_sin", mutates_args=())
>>> def numpy_sin(x: Tensor) -> Tensor:
>>>     x_np = x.cpu().numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np).to(device=x.device)
>>>
>>> x = torch.randn(3)
>>> y = numpy_sin(x)
>>> assert torch.allclose(y, x.sin())
>>>
>>> # Example of a custom op that only works for one device type.
>>> @custom_op("mylib::numpy_sin_cpu", mutates_args=(), device_types="cpu")
>>> def numpy_sin_cpu(x: Tensor) -> Tensor:
>>>     x_np = x.numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np)
>>>
>>> x = torch.randn(3)
>>> y = numpy_sin_cpu(x)
>>> assert torch.allclose(y, x.sin())
>>>
>>> # Example of a custom op that mutates an input
>>> @custom_op("mylib::numpy_sin_inplace", mutates_args={"x"}, device_types="cpu")
>>> def numpy_sin_inplace(x: Tensor) -> None:
>>>     x_np = x.numpy()
>>>     np.sin(x_np, out=x_np)
>>>
>>> x = torch.randn(3)
>>> expected = x.sin()
>>> numpy_sin_inplace(x)
>>> assert torch.allclose(x, expected)
>>>
>>> # Example of a factory function
>>> @torch.library.custom_op("mylib::bar", mutates_args={}, device_types="cpu")
>>> def bar(device: torch.device) -> Tensor:
>>>     return torch.ones(3)
>>>
>>> bar("cpu")

torch.library.triton_op(name, fn=None, /, *, mutates_args, schema=None)

创建一个自定义算子，其实现由 1 个或多个 Triton 内核支持。

这是一种更结构化的使用Triton内核与PyTorch的方式。 建议使用无需torch.library自定义操作符包装器的Triton内核 （如torch.library.custom_op()，torch.library.triton_op()），因为这样更简单； 仅在需要创建一个行为类似于PyTorch内置操作符的操作符时使用torch.library.custom_op()/torch.library.triton_op()。 例如，你可以使用一个torch.library包装API来定义当传递子类张量或在TorchDispatchMode下时Triton内核的行为。

使用 torch.library.triton_op() 而不是 torch.library.custom_op() 当实现由 1 个以上的 Triton 内核组成时。torch.library.custom_op() 将自定义算子视为不透明的（torch.compile() 和 torch.export.export() 永远不会进入它们进行追踪），但 triton_op 使实现对这些子系统的可见，允许它们优化 Triton 内核。

注意，fn 必须只包含 PyTorch 理解的操作符和 Triton 内核调用。任何在 fn 中调用的 Triton 内核都必须被包裹在一个对 torch._library.wrap_triton`() 的调用中。

Parameters

• 名称 (str) – 自定义操作的名称，格式类似于“{命名空间}::{名称}”， 例如 “mylib::my_linear”。该名称在 PyTorch 子系统（如 torch.export、FX 图）中用作操作的稳定标识符。 为避免名称冲突，请使用您的项目名称作为命名空间； 例如，pytorch/fbgemm 中的所有自定义操作都使用 “fbgemm” 作为命名空间。

• mutates_args (可迭代对象[str] 或 "unknown") – 函数变异的参数名称。 这必须准确，否则行为将无法预测。如果为“unknown”，则悲观地假设操作符的所有输入都被变异。

• schema (无 | 字符串) – 操作符的模式字符串。如果为None （推荐），我们将从其类型注释中推断出操作符的模式。除非您有特定原因，否则我们建议让系统自动推断模式。 示例：“(Tensor x, int y) -> (Tensor, Tensor)”。

Return type

可调用的

Example:

>>> import torch
>>> from torch._library import triton_op, wrap_triton
>>>
>>> import triton
>>> from triton import language as tl
>>>
>>> @triton.jit
>>> def add_kernel(
>>>     in_ptr0,
>>>     in_ptr1,
>>>     out_ptr,
>>>     n_elements,
>>>     BLOCK_SIZE: "tl.constexpr",
>>> ):
>>>     pid = tl.program_id(axis=0)
>>>     block_start = pid * BLOCK_SIZE
>>>     offsets = block_start + tl.arange(0, BLOCK_SIZE)
>>>     mask = offsets < n_elements
>>>     x = tl.load(in_ptr0 + offsets, mask=mask)
>>>     y = tl.load(in_ptr1 + offsets, mask=mask)
>>>     output = x + y
>>>     tl.store(out_ptr + offsets, output, mask=mask)
>>>
>>> @triton_op("mylib::add", mutates_args={})
>>> def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
>>>     output = torch.empty_like(x)
>>>     n_elements = output.numel()
>>>
>>>     def grid(meta):
>>>         return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
>>>
>>>     # NB: we need to wrap the triton kernel in a call to wrap_triton
>>>     wrap_triton(add_kernel)[grid](x, y, output, n_elements, 16)
>>>     return output
>>>
>>> @torch.compile
>>> def f(x, y):
>>>     return add(x, y)
>>>
>>> x = torch.randn(3, device="cuda")
>>> y = torch.randn(3, device="cuda")
>>>
>>> z = f(x, y)
>>> assert torch.allclose(z, x + y)

torch.library.wrap_triton(triton_kernel, /)

通过make_fx或非严格的torch.export捕获triton内核到图中。

这些技术执行基于调度器的跟踪（通过 __torch_dispatch__），无法看到对原始 Triton 内核的调用。 wrap_triton API 将一个 Triton 内核封装成一个可调用对象， 该对象实际上可以被追踪到图中。

请与此API一起使用torch.library.triton_op()。

示例

>>> import torch
>>> import triton
>>> from triton import language as tl
>>> from torch.fx.experimental.proxy_tensor import make_fx
>>> from torch.library import wrap_triton
>>>
>>> @triton.jit
>>> def add_kernel(
>>>     in_ptr0,
>>>     in_ptr1,
>>>     out_ptr,
>>>     n_elements,
>>>     BLOCK_SIZE: "tl.constexpr",
>>> ):
>>>     pid = tl.program_id(axis=0)
>>>     block_start = pid * BLOCK_SIZE
>>>     offsets = block_start + tl.arange(0, BLOCK_SIZE)
>>>     mask = offsets < n_elements
>>>     x = tl.load(in_ptr0 + offsets, mask=mask)
>>>     y = tl.load(in_ptr1 + offsets, mask=mask)
>>>     output = x + y
>>>     tl.store(out_ptr + offsets, output, mask=mask)
>>>
>>> def add(x, y):
>>>     output = torch.empty_like(x)
>>>     n_elements = output.numel()
>>>
>>>     def grid_fn(meta):
>>>         return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
>>>
>>>     wrap_triton(add_kernel)[grid_fn](x, y, output, n_elements, 16)
>>>     return output
>>>
>>> x = torch.randn(3, device="cuda")
>>> y = torch.randn(3, device="cuda")
>>> gm = make_fx(add)(x, y)
>>> print(gm.code)
>>> # def forward(self, x_1, y_1):
>>> #     empty_like = torch.ops.aten.empty_like.default(x_1, pin_memory = False)
>>> #     triton_kernel_wrapper_mutation_proxy = triton_kernel_wrapper_mutation(
>>> #         kernel_idx = 0, constant_args_idx = 0,
>>> #         grid = [(1, 1, 1)], kwargs = {
>>> #             'in_ptr0': x_1, 'in_ptr1': y_1, 'out_ptr': empty_like,
>>> #             'n_elements': 3, 'BLOCK_SIZE': 16
>>> #         })
>>> #     return empty_like

Return type

任何

扩展自定义操作（由Python或C++创建）

使用注册方法，如torch.library.register_kernel()和 torch.library.register_fake()，为任何操作符添加实现 （它们可能是使用torch.library.custom_op()创建的，或者通过PyTorch的C++操作符注册API进行注册）。

torch.library.register_kernel(op, device_types, func=None, /, *, lib=None)

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

一些有效的设备类型有：“cpu”，“cuda”，“xla”，“mps”，“ipu”，“xpu”。 此 API 可用作装饰器。

Parameters

• 操作 (字符串 | OpOverload) – 要注册实现的操作符。

• 设备类型 (None | str | Sequence[str]) – 要注册实现的设备类型。 如果为 None，则将注册到所有设备类型。请仅在您的实现真正与设备类型无关时使用此选项。

• 函数 (可调用对象) – 要注册为给定设备类型的实现的函数。

• 库 (可选[库]) – 如果提供，此注册的生命周期

Examples::

>>> import torch
>>> from torch import Tensor
>>> from torch.library import custom_op
>>> import numpy as np
>>>
>>> # Create a custom op that works on cpu
>>> @custom_op("mylib::numpy_sin", mutates_args=(), device_types="cpu")
>>> def numpy_sin(x: Tensor) -> Tensor:
>>>     x_np = x.numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np)
>>>
>>> # Add implementations for the cuda device
>>> @torch.library.register_kernel("mylib::numpy_sin", "cuda")
>>> def _(x):
>>>     x_np = x.cpu().numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np).to(device=x.device)
>>>
>>> x_cpu = torch.randn(3)
>>> x_cuda = x_cpu.cuda()
>>> assert torch.allclose(numpy_sin(x_cpu), x_cpu.sin())
>>> assert torch.allclose(numpy_sin(x_cuda), x_cuda.sin())

torch.library.register_autograd(op, backward, /, *, setup_context=None, lib=None)

为该自定义操作注册反向公式。

为了使算子与自动求导一起工作，你需要注册一个反向传播公式： 1. 在反向传播过程中，你需要告诉我们如何计算梯度，为此需要提供一个“反向”函数。 2. 如果你需要在前向传播中保存任何值以计算梯度，可以使用setup_context来保存这些值用于反向传播。

backward 在反向传播过程中运行。它接受 (ctx, *grads)： - grads 是一个或多个梯度。梯度的数量与操作符的输出数量相匹配。 ctx 对象是相同的 ctx 对象，由 torch.autograd.Function 使用。 backward_fn 的语义与 torch.autograd.Function.backward() 相同。

setup_context(ctx, inputs, output) 在前向传递期间运行。 请通过将所需的量保存到 ctx 对象中来准备反向传播， 可以通过torch.autograd.function.FunctionCtx.save_for_backward() 或者将其作为 ctx 的属性进行赋值。如果您的自定义操作具有仅限关键字参数，我们期望 setup_context 的签名是 setup_context(ctx, inputs, keyword_only_inputs, output)。

既setup_context_fn和backward_fn都必须是可追踪的。也就是说， 它们不能直接访问 torch.Tensor.data_ptr()，并且它们不能依赖或改变全局状态。如果你需要一个不可追踪的反向传播， 你可以将其作为一个单独的自定义操作，在backward_fn内部调用。

如果你需要在不同的设备上实现不同的自动求导行为，那么我们建议为每个需要不同行为的设备创建两个不同的自定义算子，并在运行时在这两者之间进行切换。

示例

>>> import torch
>>> import numpy as np
>>> from torch import Tensor
>>>
>>> @torch.library.custom_op("mylib::numpy_sin", mutates_args=())
>>> def numpy_sin(x: Tensor) -> Tensor:
>>>     x_np = x.cpu().numpy()
>>>     y_np = np.sin(x_np)
>>>     return torch.from_numpy(y_np).to(device=x.device)
>>>
>>> def setup_context(ctx, inputs, output) -> Tensor:
>>>     x, = inputs
>>>     ctx.save_for_backward(x)
>>>
>>> def backward(ctx, grad):
>>>     x, = ctx.saved_tensors
>>>     return grad * x.cos()
>>>
>>> torch.library.register_autograd(
...     "mylib::numpy_sin", backward, setup_context=setup_context
... )
>>>
>>> x = torch.randn(3, requires_grad=True)
>>> y = numpy_sin(x)
>>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y))
>>> assert torch.allclose(grad_x, x.cos())
>>>
>>> # Example with a keyword-only arg
>>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
>>> def numpy_mul(x: Tensor, *, val: float) -> Tensor:
>>>     x_np = x.cpu().numpy()
>>>     y_np = x_np * val
>>>     return torch.from_numpy(y_np).to(device=x.device)
>>>
>>> def setup_context(ctx, inputs, keyword_only_inputs, output) -> Tensor:
>>>     ctx.val = keyword_only_inputs["val"]
>>>
>>> def backward(ctx, grad):
>>>     return grad * ctx.val
>>>
>>> torch.library.register_autograd(
...     "mylib::numpy_mul", backward, setup_context=setup_context
... )
>>>
>>> x = torch.randn(3, requires_grad=True)
>>> y = numpy_mul(x, val=3.14)
>>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y))
>>> assert torch.allclose(grad_x, torch.full_like(x, 3.14))

torch.library.register_fake(op, func=None, /, *, lib=None, _stacklevel=1)

为该操作注册一个 FakeTensor 实现（“伪实现”）。

有时也被称为“元核”或“抽象实现”。

“‘FakeTensor 实现’”指定了此操作在不携带数据的张量（“FakeTensor”）上的行为。给定一些具有特定属性（尺寸/步幅/存储偏移/设备）的输入张量，它指定了输出张量的属性是什么。

FakeTensor 的实现与操作符具有相同的签名。 它既适用于 FakeTensors 也适用于元张量。要编写 FakeTensor 实现， 假设所有传递给操作符的张量输入都是常规的 CPU/CUDA/元张量， 但它们没有存储，并且你正试图返回常规的 CPU/CUDA/元张量作为输出。 FakeTensor 的实现必须仅由 PyTorch 操作组成（并且不能直接访问任何输入或中间张量的存储或数据）。

此 API 可用作装饰器（参见示例）。

有关自定义操作的详细指南，请参见 https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html

示例

>>> import torch
>>> import numpy as np
>>> from torch import Tensor
>>>
>>> # Example 1: an operator without data-dependent output shape
>>> @torch.library.custom_op("mylib::custom_linear", mutates_args=())
>>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor:
>>>     raise NotImplementedError("Implementation goes here")
>>>
>>> @torch.library.register_fake("mylib::custom_linear")
>>> def _(x, weight, bias):
>>>     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
>>>
>>> with torch._subclasses.fake_tensor.FakeTensorMode():
>>>     x = torch.randn(2, 3)
>>>     w = torch.randn(3, 3)
>>>     b = torch.randn(3)
>>>     y = torch.ops.mylib.custom_linear(x, w, b)
>>>
>>> assert y.shape == (2, 3)
>>>
>>> # Example 2: an operator with data-dependent output shape
>>> @torch.library.custom_op("mylib::custom_nonzero", mutates_args=())
>>> def custom_nonzero(x: Tensor) -> Tensor:
>>>     x_np = x.numpy(force=True)
>>>     res = np.stack(np.nonzero(x_np), axis=1)
>>>     return torch.tensor(res, device=x.device)
>>>
>>> @torch.library.register_fake("mylib::custom_nonzero")
>>> def _(x):
>>> # Number of nonzero-elements is data-dependent.
>>> # Since we cannot peek at the data in an fake impl,
>>> # we use the ctx object to construct a new symint that
>>> # represents the data-dependent size.
>>>     ctx = torch.library.get_ctx()
>>>     nnz = ctx.new_dynamic_size()
>>>     shape = [nnz, x.dim()]
>>>     result = x.new_empty(shape, dtype=torch.int64)
>>>     return result
>>>
>>> from torch.fx.experimental.proxy_tensor import make_fx
>>>
>>> x = torch.tensor([0, 1, 2, 3, 4, 0])
>>> trace = make_fx(torch.ops.mylib.custom_nonzero, tracing_mode="symbolic")(x)
>>> trace.print_readable()
>>>
>>> assert torch.allclose(trace(x), torch.ops.mylib.custom_nonzero(x))

torch.library.register_vmap(op, func=None, /, *, lib=None)

为该自定义操作注册一个vmap实现以支持torch.vmap()。

此 API 可用作装饰器（参见示例）。

为了使算子与torch.vmap()一起工作，您可能需要注册以下签名的vmap实现：

vmap_func(info, in_dims: Tuple[Optional[int]], *args, **kwargs),

其中 *args 和 **kwargs 是 op 的参数和关键字参数。 我们不支持仅限关键字参数的张量。

它指定了如何计算具有额外维度（由in_dims指定）的输入的op的批量版本。

对于每个args中的参数，in_dims都有一个对应的Optional[int]。如果参数不是一个张量或者参数没有被vmapped覆盖，则它是None；否则，它是一个指定张量在哪个维度上被vmapped覆盖的整数。

info 是一组可能有帮助的额外元数据： info.batch_size 指定被 vmapped 的维度的大小，而 info.randomness 是传递给 torch.vmap() 的 randomness 选项。

该函数func的返回值是一个元组，其中包含(output, out_dims)。 类似于in_dims， out_dims应该具有相同的结构，并且每个输出中都包含一个out_dim， 以指定输出是否具有vmapped维度及其索引。

示例

>>> import torch
>>> import numpy as np
>>> from torch import Tensor
>>> from typing import Tuple
>>>
>>> def to_numpy(tensor):
>>>     return tensor.cpu().numpy()
>>>
>>> lib = torch.library.Library("mylib", "FRAGMENT")
>>> @torch.library.custom_op("mylib::numpy_cube", mutates_args=())
>>> def numpy_cube(x: Tensor) -> Tuple[Tensor, Tensor]:
>>>     x_np = to_numpy(x)
>>>     dx = torch.tensor(3 * x_np ** 2, device=x.device)
>>>     return torch.tensor(x_np ** 3, device=x.device), dx
>>>
>>> def numpy_cube_vmap(info, in_dims, x):
>>>     result = numpy_cube(x)
>>>     return result, (in_dims[0], in_dims[0])
>>>
>>> torch.library.register_vmap(numpy_cube, numpy_cube_vmap)
>>>
>>> x = torch.randn(3)
>>> torch.vmap(numpy_cube)(x)
>>>
>>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
>>> def numpy_mul(x: Tensor, y: Tensor) -> Tensor:
>>>     return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device)
>>>
>>> @torch.library.register_vmap("mylib::numpy_mul")
>>> def numpy_mul_vmap(info, in_dims, x, y):
>>>     x_bdim, y_bdim = in_dims
>>>     x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1)
>>>     y = y.movedim(y_bdim, -1) if y_bdim is not None else y.unsqueeze(-1)
>>>     result = x * y
>>>     result = result.movedim(-1, 0)
>>>     return result, 0
>>>
>>>
>>> x = torch.randn(3)
>>> y = torch.randn(3)
>>> torch.vmap(numpy_mul)(x, y)

注意

vmap函数应旨在保留整个自定义算子的语义。 也就是说，grad(vmap(op)) 应该可以被 grad(map(op)) 替代。

如果你的自定义算子在反向传递中有任何自定义行为，请记住这一点。

torch.library.impl_abstract(qualname, func=None, *, lib=None, _stacklevel=1)

该API在PyTorch 2.4中重命名为torch.library.register_fake()。 请使用那个代替。

torch.library.get_ctx()

get_ctx() 返回当前的 AbstractImplCtx 对象。

调用 get_ctx() 只能在 fake impl 内有效 （有关更多用法细节，请参见 torch.library.register_fake()）。

Return type

FakeImplCtx

torch.library.register_torch_dispatch(op, torch_dispatch_class, func=None, /, *, lib=None)

为给定的操作符注册一个torch_dispatch规则和torch_dispatch_class。

这允许开放注册来指定操作符和torch_dispatch_class之间的行为，而无需修改torch_dispatch_class或操作符本身。

torch_dispatch_class 是具有 __torch_dispatch__ 的 Tensor 子类或是一个 TorchDispatchMode。

如果是一个张量子类，我们期望 func 具有以下签名： (cls, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any

如果它是TorchDispatchMode，我们期望func具有以下签名： (mode, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any

args 和 kwargs 将以相同的方式进行归一化，就像在 __torch_dispatch__ 中一样（参见 __torch_dispatch__ 调用约定）。

示例

>>> import torch
>>>
>>> @torch.library.custom_op("mylib::foo", mutates_args={})
>>> def foo(x: torch.Tensor) -> torch.Tensor:
>>>     return x.clone()
>>>
>>> class MyMode(torch.utils._python_dispatch.TorchDispatchMode):
>>>     def __torch_dispatch__(self, func, types, args=(), kwargs=None):
>>>         return func(*args, **kwargs)
>>>
>>> @torch.library.register_torch_dispatch("mylib::foo", MyMode)
>>> def _(mode, func, types, args, kwargs):
>>>     x, = args
>>>     return x + 1
>>>
>>> x = torch.randn(3)
>>> y = foo(x)
>>> assert torch.allclose(y, x)
>>>
>>> with MyMode():
>>>     y = foo(x)
>>> assert torch.allclose(y, x + 1)

torch.library.infer_schema(prototype_function, /, *, mutates_args, op_name=None)

解析给定函数的模式，该模式带有类型提示。模式是从函数的类型提示中推断出来的，并可用于定义一个新的操作符。

我们做出如下假设：

• 没有任何输出与输入或彼此重复。

• 未指定库的字符串类型注解“设备、数据类型、张量、类型”

  假定为 torch.*。同样地，“Optional、List、Sequence、Union”字符串类型注解。

  未指定库的情况均被视为typing.*。
• 只有mutates_args中列出的参数会被修改。如果mutates_args是“未知”，

  它假设该操作的所有输入都在被修改。

调用者（例如自定义操作API）负责检查这些假设。

Parameters

• 原型函数 (可调用对象) – 从其类型注释推断模式的函数。

• op_name (可选[字符串]) – 模式中操作符的名称。如果 name 为 None，则名称不会包含在推断的模式中。请注意，torch.library.Library.define 的输入模式需要一个操作符名称。

• mutates_args ("未知" | 可迭代对象[str]) – 函数中被修改的参数。

Returns

推断的模式。

Return type

字符串

示例

>>> def foo_impl(x: torch.Tensor) -> torch.Tensor:
>>>     return x.sin()
>>>
>>> infer_schema(foo_impl, op_name="foo", mutates_args={})
foo(Tensor x) -> Tensor
>>>
>>> infer_schema(foo_impl, mutates_args={})
(Tensor x) -> Tensor

class torch._library.custom_ops.CustomOpDef(namespace, name, schema, fn)

CustomOpDef 是一个围绕函数创建自定义操作的包装器。

它提供了多种方法来为此自定义操作注册附加行为。

你不应直接实例化CustomOpDef；相反，使用 torch.library.custom_op() API。

set_kernel_enabled(device_type, enabled=True)

禁用或重新启用此自定义算子已注册的内核。

如果内核已经是禁用/启用状态，这将不起作用。

注意

如果一个内核首先被禁用然后再注册，它会一直保持禁用状态直到再次启用。

Parameters

• 设备类型 (str) – 禁用/启用内核的设备类型。

• 禁用 (bool) – 是否禁用或启用内核。

示例

>>> inp = torch.randn(1)
>>>
>>> # define custom op `f`.
>>> @custom_op("mylib::f", mutates_args=())
>>> def f(x: Tensor) -> Tensor:
>>>     return torch.zeros(1)
>>>
>>> print(f(inp))  # tensor([0.]), default kernel
>>>
>>> @f.register_kernel("cpu")
>>> def _(x):
>>>     return torch.ones(1)
>>>
>>> print(f(inp))  # tensor([1.]), CPU kernel
>>>
>>> # temporarily disable the CPU kernel
>>> with f.set_kernel_enabled("cpu", enabled = False):
>>>     print(f(inp))  # tensor([0.]) with CPU kernel disabled

低级APIs

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

警告

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

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

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

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

要创建一个用于重载现有库（名称为 ns）操作符的库，请将类型设置为“IMPL”。 要创建一个新的库（名称为 ns）以注册新的操作符，请将类型设置为“DEF”。 要创建一个可能已存在的库片段以注册操作符（并绕过给定命名空间只能有一个库的限制），请将类型设置为“FRAGMENT”。

Parameters

• ns – 库名称

• 类型 – “DEF”，“IMPL”（默认：“IMPL”），“FRAGMENT”

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

define(schema, alias_analysis='', *, tags=())

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

Parameters

• 模式 – 定义新算子的功能模式。

• alias分析 (可选) – 表示操作数的别名属性是否可以从模式（默认行为）中推断出来，或者不能（“保守”）。

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

Returns

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

Example::

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

fallback(fn, dispatch_key='', *, with_keyset=False)

将该函数实现注册为给定键的备用选项。

此功能仅适用于具有全局命名空间的库 (“_”)。

Parameters

• fn – 用作给定分派键的备用函数，或fallthrough_kernel() 以注册一个贯穿。

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

• with_keyset – 控制当前调度器调用的键集是否应作为第一个参数传递给调用时的 fn。这应该用于为重新调度调用创建适当的键集。

Example::

>>> my_lib = Library("_", "IMPL")
>>> def fallback_kernel(op, *args, **kwargs):
>>>     # Handle all autocast ops generically
>>>     # ...
>>> my_lib.fallback(fallback_kernel, "Autocast")

impl(op_name, fn, dispatch_key='', *, with_keyset=False)

为库中定义的操作符注册函数实现。

Parameters

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

• fn – 输入分发键的操作实现函数，或fallthrough_kernel() 以注册一个备用操作。

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

• with_keyset – 控制当前调度器调用的键集是否应作为第一个参数传递给调用时的 fn。这应该用于为重新调度调用创建适当的键集。

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以注册一个贯穿。

torch.library.define(qualname, schema, *, lib=None, tags=())

torch.library.define(lib, schema, alias_analysis='')

定义一个新的算子。

在 PyTorch 中，定义一个算子（op 的全称为“operator”）是一个两步过程： - 我们需要定义该算子（提供算子名称和模式） - 我们需要实现该算子与各种 PyTorch 子系统的交互行为，比如 CPU/CUDA 张量、自动微分等。

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

Parameters

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

• 模式 (字符串) – 算子的模式。例如，对于接受一个张量并返回一个张量的操作，其模式为 “(Tensor x) -> Tensor”。 这里不包含操作名称（操作名称通过 qualname 传递）。

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

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

Example::

>>> import torch
>>> import numpy as np
>>>
>>> # Define the operator
>>> torch.library.define("mylib::sin", "(Tensor x) -> Tensor")
>>>
>>> # Add implementations for the operator
>>> @torch.library.impl("mylib::sin", "cpu")
>>> def f(x):
>>>     return torch.from_numpy(np.sin(x.numpy()))
>>>
>>> # Call the new operator from torch.ops.
>>> x = torch.randn(3)
>>> y = torch.ops.mylib.sin(x)
>>> assert torch.allclose(y, x.sin())

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) – 应该是一个类似“命名空间::操作符名称”的字符串。

• 类型 (str | Sequence[str]) – 要注册实现的设备类型。

• 库 (可选[库]) – 如果提供，此注册的生命周期将与库对象的生命周期相关联。

示例

>>> import torch
>>> import numpy as np
>>>
>>> # Define the operator
>>> torch.library.define("mylib::mysin", "(Tensor x) -> Tensor")
>>>
>>> # Add implementations for the cpu device
>>> @torch.library.impl("mylib::mysin", "cpu")
>>> def f(x):
>>>     return torch.from_numpy(np.sin(x.numpy()))
>>>
>>> x = torch.randn(3)
>>> y = torch.ops.mylib.mysin(x)
>>> assert torch.allclose(y, x.sin())