UX 限制¶

torch.func，类似于 JAX，对可以转换的内容有限制。通常，JAX 的限制是转换仅适用于纯函数：也就是说，输出完全由输入决定的函数，并且不涉及副作用（如变异）。

我们有类似的保证：我们的转换适用于纯函数。但是，我们确实支持某些就地操作。一方面，编写与函数转换兼容的代码可能涉及更改您编写 PyTorch 代码的方式，另一方面，您可能会发现我们的转换可以让您表达以前在 PyTorch 中难以表达的东西。

一般限制¶

所有 torch.func 转换都共享一个限制，即函数不应该分配给全局变量。相反，函数的所有输出必须从函数中返回。此限制来自 torch.func 的实现方式：每个转换都将张量输入包装在特殊的 torch.func 张量子类中，这些子类有助于转换。

因此，不要使用以下方法

import torch
from torch.func import grad

# Don't do this
intermediate = None

def f(x):
  global intermediate
  intermediate = x.sin()
  z = intermediate.sin()
  return z

x = torch.randn([])
grad_x = grad(f)(x)

请将 f 重写为返回 intermediate

def f(x):
  intermediate = x.sin()
  z = intermediate.sin()
  return z, intermediate

grad_x, intermediate = grad(f, has_aux=True)(x)

torch.autograd API¶

如果您尝试使用 torch.autograd API（如 torch.autograd.grad 或 torch.autograd.backward）在由 vmap() 或 torch.func 的 AD 转换（vjp()、jvp()、jacrev()、jacfwd()）转换的函数内部，则转换可能无法对其进行转换。如果无法转换，您将收到错误消息。

这是 PyTorch 的 AD 支持实现方式中的一个基本设计限制，也是我们设计 torch.func 库的原因。请改用 torch.autograd API 的 torch.func 等效项： - torch.autograd.grad、Tensor.backward -> torch.func.vjp 或 torch.func.grad - torch.autograd.functional.jvp -> torch.func.jvp - torch.autograd.functional.jacobian -> torch.func.jacrev 或 torch.func.jacfwd - torch.autograd.functional.hessian -> torch.func.hessian

vmap 限制¶

注意

vmap() 是我们限制最严格的转换。与梯度相关的转换（grad()、vjp()、jvp()）没有这些限制。 jacfwd()（以及 hessian()，它是使用 jacfwd() 实现的）是 vmap() 和 jvp() 的组合，因此它也具有这些限制。

vmap(func) 是一个转换，它返回一个函数，该函数将 func 映射到每个输入张量的新维度上。vmap 的心理模型是它就像运行一个 for 循环：对于纯函数（即在没有副作用的情况下），vmap(f)(x) 等效于

torch.stack([f(x_i) for x_i in x.unbind(0)])

变异：Python 数据结构的任意变异¶

在存在副作用的情况下，vmap()不再像运行一个for循环那样工作。例如，以下函数

def f(x, list):
  list.pop()
  print("hello!")
  return x.sum(0)

x = torch.randn(3, 1)
lst = [0, 1, 2, 3]

result = vmap(f, in_dims=(0, None))(x, lst)

将只打印一次“hello！”并且只从 lst 中弹出单个元素。

vmap() 只执行一次 f，因此所有副作用都只发生一次。

这是 vmap 实现方式的结果。torch.func 有一个特殊的内部 BatchedTensor 类。 vmap(f)(*inputs) 获取所有 Tensor 输入，将它们转换为 BatchedTensors，并调用 f(*batched_tensor_inputs)。BatchedTensor 覆盖 PyTorch API 以对每个 PyTorch 运算符产生批处理（即矢量化）行为。

变异：就地 PyTorch 运算¶

您可能是因为收到有关 vmap 不兼容的就地运算的错误而来到这里的。 vmap() 如果遇到不支持的 PyTorch 就地运算，就会引发错误，否则会成功。不支持的运算会将具有更多元素的 Tensor 写入具有较少元素的 Tensor。以下是如何发生这种情况的示例

def f(x, y):
  x.add_(y)
  return x

x = torch.randn(1)
y = torch.randn(3, 1)  # When vmapped over, looks like it has shape [1]

# Raises an error because `x` has fewer elements than `y`.
vmap(f, in_dims=(None, 0))(x, y)

x 是一个具有一个元素的 Tensor，y 是一个具有三个元素的 Tensor。 x + y 具有三个元素（由于广播），但尝试将三个元素写回 x，而它只有一个元素，则由于试图将三个元素写入一个只有一个元素的 Tensor，因此会引发错误。

如果在 vmap() 下被批处理的 Tensor（即，它正在被 vmap），则不会出现问题。

def f(x, y):
  x.add_(y)
  return x

x = torch.randn(3, 1)
y = torch.randn(3, 1)
expected = x + y

# Does not raise an error because x is being vmapped over.
vmap(f, in_dims=(0, 0))(x, y)
assert torch.allclose(x, expected)

一个常见的解决方法是将对工厂函数的调用替换为其“new_*”等效项。例如

用 Tensor.new_zeros() 替换 torch.zeros()
用 Tensor.new_empty() 替换 torch.empty()

要了解为什么这会有所帮助，请考虑以下内容。

def diag_embed(vec):
  assert vec.dim() == 1
  result = torch.zeros(vec.shape[0], vec.shape[0])
  result.diagonal().copy_(vec)
  return result

vecs = torch.tensor([[0., 1, 2], [3., 4, 5]])

# RuntimeError: vmap: inplace arithmetic(self, *extra_args) is not possible ...
vmap(diag_embed)(vecs)

在 vmap() 内部，result 是一个形状为 [3, 3] 的 Tensor。但是，虽然 vec 看起来形状为 [3]，但 vec 实际上具有底层形状 [2, 3]。无法将 vec 复制到 result.diagonal() 中，因为它的形状是 [3]，因为它具有太多元素。

def diag_embed(vec):
  assert vec.dim() == 1
  result = vec.new_zeros(vec.shape[0], vec.shape[0])
  result.diagonal().copy_(vec)
  return result

vecs = torch.tensor([[0., 1, 2], [3., 4, 5]])
vmap(diag_embed)(vecs)

用 Tensor.new_zeros() 替换 torch.zeros() 使得 result 具有形状为 [2, 3, 3] 的底层 Tensor，因此现在可以将 vec（其底层形状为 [2, 3]）复制到 result.diagonal() 中。

变异：out= PyTorch 运算¶

vmap() 不支持 PyTorch 运算中的 out= 关键字参数。如果在您的代码中遇到它，它会友好地出错。

这不是一个根本的限制；从理论上讲，我们将来可以支持它，但目前我们选择不这样做。

数据依赖的 Python 控制流¶

我们目前不支持 vmap 对数据依赖的控制流进行转换。数据依赖的控制流是指当 if 语句、while 循环或 for 循环的条件是一个被 vmap 处理的 Tensor 时的情况。例如，以下操作将引发错误消息

def relu(x):
  if x > 0:
    return x
  return 0

x = torch.randn(3)
vmap(relu)(x)

但是，任何不依赖于 vmap 处理的 Tensor 中的值的控制流都可以正常工作

def custom_dot(x):
  if x.dim() == 1:
    return torch.dot(x, x)
  return (x * x).sum()

x = torch.randn(3)
vmap(custom_dot)(x)

JAX 支持使用特殊的控制流运算符（例如 jax.lax.cond、jax.lax.while_loop）对数据依赖的控制流进行转换。我们正在研究将这些运算符的等效项添加到 PyTorch 中。

数据依赖的运算 (.item())¶

我们不（也不打算）支持对调用 .item() 来操作 Tensor 的用户定义函数进行 vmap。例如，以下操作将引发错误消息

def f(x):
  return x.item()

x = torch.randn(3)
vmap(f)(x)

请尝试重写您的代码，避免使用 .item() 调用。

您也可能会遇到有关使用 .item() 的错误消息，但您可能没有使用它。在这些情况下，PyTorch 内部可能正在调用 .item() - 请在 GitHub 上提交问题，我们会修复 PyTorch 内部。

动态形状运算（nonzero 及其相关运算）¶

vmap(f) 要求将 f 应用于输入中的每个“示例”都返回具有相同形状的 Tensor。不支持 torch.nonzero、torch.is_nonzero 等运算，因此它们会引发错误。

要了解原因，请考虑以下示例

xs = torch.tensor([[0, 1, 2], [0, 0, 3]])
vmap(torch.nonzero)(xs)

torch.nonzero(xs[0]) 返回一个形状为 2 的 Tensor；但 torch.nonzero(xs[1]) 返回一个形状为 1 的 Tensor。我们无法构造一个单个 Tensor 作为输出；输出需要是一个不规则 Tensor（而 PyTorch 还没有不规则 Tensor 的概念）。

随机性¶

用户在调用随机运算时的意图可能不清楚。具体来说，一些用户可能希望在整个批次中随机行为保持一致，而另一些用户则希望随机行为在整个批次中不同。为了解决这个问题，vmap 使用一个随机性标志。

该标志只能传递给 vmap，并且可以取三个值：“error”、“different” 或 “same”，默认值为 error。在 “error” 模式下，对随机函数的任何调用都会产生一个错误，要求用户根据其使用情况使用另外两个标志中的一个。

在 “different” 随机性下，批次中的元素会产生不同的随机值。例如，

def add_noise(x):
  y = torch.randn(())  # y will be different across the batch
  return x + y

x = torch.ones(3)
result = vmap(add_noise, randomness="different")(x)  # we get 3 different values

在 “same” 随机性下，批次中的元素会产生相同的随机值。例如，

def add_noise(x):
  y = torch.randn(())  # y will be the same across the batch
  return x + y

x = torch.ones(3)
result = vmap(add_noise, randomness="same")(x)  # we get the same value, repeated 3 times

警告

我们的系统只能确定 PyTorch 运算符的随机性行为，无法控制其他库（如 numpy）的行为。这类似于 JAX 解决方案中的限制

注意

使用任何一种支持的随机性进行多次 vmap 调用不会产生相同的结果。与标准 PyTorch 一样，用户可以通过在 vmap 之外使用 torch.manual_seed() 或使用生成器来获得随机性可重复性。

注意

最后，我们的随机性与 JAX 不同，因为我们没有使用无状态 PRNG，部分原因是 PyTorch 不完全支持无状态 PRNG。相反，我们引入了一个标志系统，以允许我们看到的随机性的最常见形式。如果您的使用情况不适合这些形式的随机性，请提交问题。