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 中的每個 arg，in_dims 有一個相應的 Optional[int]。如果 arg 不是 Tensor 或 arg 不被 vmap，則為 None，否則，它是一個指定 Tensor 的哪個維度被 vmap 的整數。

• *args，與 forward() 的 args 相同。

vmap 靜態方法的返回值是一個元組 (output, out_dims)。與 in_dims 類似，out_dims 的結構應與 output 相同，並且每個輸出都包含一個 out_dim，指定輸出是否具有 vmap 的維度以及在該維度中的索引。

有關更多詳細資訊，請參閱 使用 autograd.Function 擴充套件 torch.func。