recipes/torch_compile_torch_function_modes

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

(beta) 利用 Torch Function 模式与 torch.compile

作者： Michael Lazos

本教程介绍如何使用 PyTorch 的一个关键扩展点，

torch function 模式，与 torch.compile 结合使用，在追踪时覆盖 torch 算子（也称为 ops）的行为，且不产生运行时开销。

注意

本教程需要 PyTorch 2.7.0 或更高版本。

重写一个 torch 算子 (torch.add -> torch.mul)

对于这个例子，我们将使用 torch function 模式将加法运算重写为乘法运算。如果某个后端具有应该为给定算子分发的自定义实现，则这种覆盖可能很常见。

import torch

# exit cleanly if we are on a device that doesn't support ``torch.compile``
if torch.cuda.get_device_capability() < (7, 0):
    print("Exiting because torch.compile is not supported on this device.")
    import sys
    sys.exit(0)

from torch.overrides import BaseTorchFunctionMode

# Define our mode, Note: ``BaseTorchFunctionMode``
# implements the actual invocation of func(..)
class AddToMultiplyMode(BaseTorchFunctionMode):
    def __torch_function__(self, func, types, args=(), kwargs=None):
        if func == torch.Tensor.add:
            func = torch.mul

        return super().__torch_function__(func, types, args, kwargs)

@torch.compile()
def test_fn(x, y):
    return x + y * x # Note: infix operators map to torch.Tensor.* methods

x = torch.rand(2, 2)
y = torch.rand_like(x)

with AddToMultiplyMode():
    z = test_fn(x, y)

assert torch.allclose(z, x * y * x)

# The mode can also be used within the compiled region as well like this:

@torch.compile()
def test_fn(x, y):
    with AddToMultiplyMode():
        return x + y * x # Note: infix operators map to torch.Tensor.* methods

x = torch.rand(2, 2)
y = torch.rand_like(x)
z = test_fn(x, y)

assert torch.allclose(z, x * y * x)

结论

在本教程中，我们演示了如何使用 torch.compile 中的 torch function 模式来覆盖 torch.* 算子的行为。这使得用户能够在不产生每次算子调用运行时开销的情况下，利用 torch function 模式的扩展优势。

• 有关其他示例和 Torch Function 模式的背景信息，请参阅 使用模式扩展 Torch API。