扩展 torch 类型，使其具有类似 Tensor 的类型功能#
此功能受到 NumPy __array_function__ 协议的启发。更多细节请参阅 NumPy 文档和 NEP-0018。
为了具体说明，从简单的示例开始，该示例展示了 API 分派机制。将创建自定义类型，该类型表示 2D 标量张量，参数化由顺序 N 和对角线元素的值 value 决定。
import torch
class ScalarTensor:
def __init__(self, N, value):
self._N = N
self._value = value
def __repr__(self):
return f"ScalarTensor(N={self._N}, value={self._value})"
def tensor(self):
return self._value * torch.eye(self._N)
这个设计的第一版并没有什么用处。 ScalarTensor 的主要功能是提供比基类张量更紧凑的标量张量字符串表示形式：
d = ScalarTensor(5, 2)
d, d.tensor()
(ScalarTensor(N=5, value=2),
tensor([[2., 0., 0., 0., 0.],
[0., 2., 0., 0., 0.],
[0., 0., 2., 0., 0.],
[0., 0., 0., 2., 0.],
[0., 0., 0., 0., 2.]]))
如果尝试使用此对象与 torch API，将会遇到问题：
import torch
torch.mean(d)
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[3], line 2
1 import torch
----> 2 torch.mean(d)
TypeError: mean(): argument 'input' (position 1) must be Tensor, not ScalarTensor
向 ScalarTensor 添加 __torch_function__ 实现可以使上述作成功。重新执行，这次添加 __torch_function__ 实现：
HANDLED_FUNCTIONS = {}
class ScalarTensor:
def __init__(self, N, value):
self._N = N
self._value = value
def __repr__(self):
return "ScalarTensor(N={}, value={})".format(self._N, self._value)
def tensor(self):
return self._value * torch.eye(self._N)
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
if func not in HANDLED_FUNCTIONS or not all(
issubclass(t, (torch.Tensor, ScalarTensor))
for t in types
):
return NotImplemented
return HANDLED_FUNCTIONS[func](*args, **kwargs)
__torch_function__ 方法接受四个参数：func，对正在重写的 torch API 函数的引用，types，实现 __torch_function__ 的 Tensor 类类型列表，args，传递给函数的参数元组，以及 kwargs，传递给函数的关键字的字典。它使用名为 HANDLED_FUNCTIONS 存储自定义实现。该字典的键是 torch 命名空间中的函数，值是 ScalarTensor 的实现。
备注
使用全局 global dispatch table 不是 __torch_function__ API 部分，它只是构建覆盖实现的有用设计模式。
当传递 ScalarTensor 时，这个类定义还不足以让 torch.mean 做正确的事情——还需要为 ScalarTensor operand 定义 torch.mean 的实现，并将该实现添加到 HANDLED_FUNCTIONS dispatch 表字典中。一种方法是定义装饰器：
import functools
def implements(torch_function):
"""Register a torch function override for ScalarTensor"""
def decorator(func):
functools.update_wrapper(func, torch_function)
HANDLED_FUNCTIONS[torch_function] = func
return func
return decorator
这可以应用于覆盖的实现：
@implements(torch.mean)
def mean(input):
return float(input._value) / input._N
通过此更改，现在可以将 torch.mean() 与 ScalarTensor 一起使用：
d = ScalarTensor(5, 2)
torch.mean(d)
当然，torch.mean() 是最简单的覆盖函数的例子，因为它只需要一个 operand。可以使用相同的机制来覆盖接受多个 operand 的函数，其中任何一个都可能是定义 __torch_function__ 的张量或类似张量的函数，例如 torch.add()：
def ensure_tensor(data):
if isinstance(data, ScalarTensor):
return data.tensor()
return torch.as_tensor(data)
@implements(torch.add)
def add(input, other):
try:
if input._N == other._N:
return ScalarTensor(input._N, input._value + other._value)
else:
raise ValueError("Shape mismatch!")
except AttributeError:
return torch.add(ensure_tensor(input), ensure_tensor(other))
此版本在两个操作数均为实例时具有一条快速路径，而在任一操作数不是实例时则会退化为将数据转换为张量的较慢路径。这样处理使得当任一操作数是 ScalarTensor 或常规 Tensor 时，覆盖函数都能正确运行。
s = ScalarTensor(2, 2)
t = torch.tensor([[1, 1,], [1, 1]])
torch.add(s, s), torch.add(s, t)
请注意，add 实现不将 alpha 或 out 作为关键字参数，就像 torch.add() 那样：
torch.add(s, s, alpha=2)
为了速度和灵活性，__torch_function__ 分发机制不会检查重写函数的签名是否与 torch API 中被重写函数的签名匹配。对于某些应用来说，忽略可选参数可能没有问题，但为了确保与 Tensor 的完全兼容性，用户实现的 torch API 函数应注意精确模拟被重写函数的 API。
在 torch API 中没有显式重写的函数将从 __torch_function__ 返回 NotImplemented。如果所有定义了 __torch_function__ 的操作数都返回 NotImplemented，PyTorch 将引发 TypeError。这意味着大多数情况下，当传递该类型的实例时，没有为该类型显式重写的算子将引发TypeError。
try:
torch.mul(s, 3)
except TypeError:
...
实际上这意味着，如果希望使用类似 __torch_function__ 的实现来重载，你需要显式地实现完整的 torch API 或你使用案例中关心的 API 的整个子集。这可能是艰巨的任务，因为完整的 torch API 非常庞大。
另一个选择是，对于未处理的算子，不返回 NotImplemented，而是在没有重写可用时将 Tensor 传递给原始的 torch 函数。例如，如果将 ScalarTensor 的 __torch_function__ 实现更改为如下所示：
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
if func not in HANDLED_FUNCTIONS or not all(
issubclass(t, (torch.Tensor, ScalarTensor))
for t in types
):
args = [a.tensor() if hasattr(a, 'tensor') else a for a in args]
return func(*args, **kwargs)
return HANDLED_FUNCTIONS[func](*args, **kwargs)
那么 torch.mul() 就能正常工作，尽管返回值类型始终是 Tensor 而不是 ScalarTensor，即使两个操作数都是 ScalarTensor 实例：
s = ScalarTensor(2, 2)
torch.mul(s, s)
另请参阅下文的 MetadataTensor 示例，展示这种模式的另一种变体，但该变体始终返回 MetadataTensor 以在 torch API 的算子中传播元数据。
__torch_function__ 协议的设计旨在全面覆盖 API，部分覆盖可能导致不良后果，特别是某些函数会引发 TypeError。对于子类尤其如此，必须同时覆盖torch.add()、torch.Tensor.__add__() 和 torch.Tensor.add() 这三个方法，即使它们返回完全相同的结果。未能做到这一点还可能导致无限递归。如果需要从 torch.Tensor 子类实现某个函数，则必须在实现中使用 super().__torch_function__。