torch.autograd.Function.vmap

静态Function.vmap(info, in_dims, *args)[源代码]

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

为了使torch.autograd.Function()支持torch.vmap()，你必须重写这个静态方法或将generate_vmap_rule设置为True（但不能同时进行）。

如果你选择覆盖这个静态方法：它必须接受

将一个info对象作为第一个参数。其中，info.batch_size指定了被vmapped的维度大小，而info.randomness是传递给torch.vmap()的随机性选项。
作为第二个参数传递一个in_dims元组。对于args中的每个参数，in_dims包含一个对应的Optional[int]值。如果参数不是张量或不被vmapped处理，则该值为None；否则，它是一个整数，表示正在被vmapped的张量的维度。
*args，这与传递给forward()方法的参数相同。

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