Pytorch使用手册（专题五十）—自定义运算符

1. PyTorch 自定义运算符

PyTorch 提供了一个庞大的运算符库，这些运算符可以对张量进行操作（例如 torch.add、torch.sum 等）。然而，您可能希望向 PyTorch 引入一个新的自定义操作，并使其能够与诸如 torch.compile、autograd 和 torch.vmap 等子系统协同工作。为此，您必须通过 Python 的 torch.library 文档或 C++ 的 TORCH_LIBRARY API 将自定义操作注册到 PyTorch 中。

1.1 从 Python 编写自定义运算符

请参阅 自定义 Python 运算符。

如果您有以下需求，您可能希望从 Python（而不是 C++）编写自定义运算符：

1. 您有一个 Python 函数，希望 PyTorch 将其视为一个不透明的可调用对象，特别是在 torch.compile 和 torch.export 方面。

2. 您有一些 Python 绑定的 C++/CUDA 内核，并希望这些内核能够与 PyTorch 子系统ÿ