torch.autograd.Function.vmap

static Function.vmap(info, in_dims, *args)

定义此 autograd.Function 在 torch.vmap() 下的行为。

要让一个 torch.autograd.Function() 支持 torch.vmap()，您必须覆盖此静态方法，或者将 generate_vmap_rule 设置为 True（您不能同时进行这两项）。

如果您选择重写此静态方法：它必须接受

• 一个 info 对象作为第一个参数。 info.batch_size 指定了 vmap 操作的维度大小，而 info.randomness 是传递给 torch.vmap() 的随机性选项。

• 一个 in_dims 元组作为第二个参数。对于 args 中的每个参数，in_dims 都有一个对应的 Optional[int]。如果参数不是 Tensor 或未被 vmap，则为 None，否则，它是一个整数，指定了 Tensor 中正在被 vmap 的维度。

• *args，与传递给 torch.autograd.Function.forward() 的参数相同。

vmap 静态方法的返回值为一个 (output, out_dims) 元组。与 in_dims 类似，out_dims 的结构应与 output 相同，并且每个输出都包含一个 out_dim，指定输出是否具有 vmap 维度以及该维度在该输出中的索引。

有关更多详细信息，请参阅 使用 autograd.Function 扩展 torch.func。