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，与 forward() 方法的参数相同。

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

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