NestedIOFunction

class torch.autograd.function.NestedIOFunction(*args, **kwargs)

此類僅出於向後相容性原因而存在。對於任何新的用例，請改用 Function。

backward(*gradients)

共享的後向工具函式。

返回型別

任意型別

backward_extended(*grad_output)

使用者定義的後向。

forward(*args)

共享的前向工具函式。

返回型別

任意型別

forward_extended(*input)

使用者定義的前向。

static jvp(ctx, *grad_inputs)

定義使用前向模式自動微分來區分操作的公式。

所有子類都必須重寫此函式。它必須接受一個上下文 ctx 作為第一個引數，後跟與 forward() 接收到的輸入一樣多的引數（對於 forward 函式的非 Tensor 輸入，將傳入 None），並且它應該返回與 forward() 輸出的 Tensor 一樣多的 Tensor。每個引數是相對於給定輸入的梯度，每個返回值應該是相對於相應輸出的梯度。如果輸出不是 Tensor 或者函式相對於該輸出不可微分，則可以簡單地為該輸入傳遞 None 作為梯度。

您可以使用 ctx 物件將 forward 中的任何值傳遞給此函式。

返回型別

任意型別

mark_dirty(*args, **kwargs)

請參閱 Function.mark_dirty()。

mark_non_differentiable(*args, **kwargs)

請參閱 Function.mark_non_differentiable()。

save_for_backward(*args)

請參閱 Function.save_for_backward()。

save_for_forward(*tensors)

儲存給定的 Tensor 以供將來呼叫 jvp()。

save_for_forward 最多隻能呼叫一次，可以在 setup_context() 或 forward() 方法中呼叫，並且所有引數都應該是 Tensor。

在 jvp() 中，可以透過 saved_tensors 屬性訪問儲存的物件。

引數也可以是 None。這沒有操作。

有關如何使用此方法的更多詳細資訊，請參閱 擴充套件 torch.autograd。

示例：

>>> class Func(torch.autograd.Function):
>>>     @staticmethod
>>>     def forward(ctx, x: torch.Tensor, y: torch.Tensor, z: int):
>>>         ctx.save_for_backward(x, y)
>>>         ctx.save_for_forward(x, y)
>>>         ctx.z = z
>>>         return x * y * z
>>>
>>>     @staticmethod
>>>     def jvp(ctx, x_t, y_t, _):
>>>         x, y = ctx.saved_tensors
>>>         z = ctx.z
>>>         return z * (y * x_t + x * y_t)
>>>
>>>     @staticmethod
>>>     def vjp(ctx, grad_out):
>>>         x, y = ctx.saved_tensors
>>>         z = ctx.z
>>>         return z * grad_out * y, z * grad_out * x, None
>>>
>>>     a = torch.tensor(1., requires_grad=True, dtype=torch.double)
>>>     t = torch.tensor(1., dtype=torch.double)
>>>     b = torch.tensor(2., requires_grad=True, dtype=torch.double)
>>>     c = 4
>>>
>>>     with fwAD.dual_level():
>>>         a_dual = fwAD.make_dual(a, t)
>>>         d = Func.apply(a_dual, b, c)

property saved_tensors

請參閱 Function.saved_tensors()。

set_materialize_grads(value)

設定是否具體化（materialize）梯度 Tensor。預設為 True。

這隻能在 setup_context() 或 forward() 方法中呼叫。

如果為 True，則在呼叫 backward() 和 jvp() 方法之前，未定義的梯度 Tensor 將被擴充套件為充滿零的 Tensor。

示例：

>>> class SimpleFunc(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         return x.clone(), x.clone()
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):
>>>         return g1 + g2  # No check for None necessary
>>>
>>> # We modify SimpleFunc to handle non-materialized grad outputs
>>> class Func(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         ctx.set_materialize_grads(False)
>>>         ctx.save_for_backward(x)
>>>         return x.clone(), x.clone()
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):
>>>         x, = ctx.saved_tensors
>>>         grad_input = torch.zeros_like(x)
>>>         if g1 is not None:  # We must check for None now
>>>             grad_input += g1
>>>         if g2 is not None:
>>>             grad_input += g2
>>>         return grad_input
>>>
>>> a = torch.tensor(1., requires_grad=True)
>>> b, _ = Func.apply(a)  # induces g2 to be undefined

static setup_context(ctx, inputs, output)

定義 autograd.Function 的前向過程有兩種方法。

或者

1. 使用簽名 forward(ctx, *args, **kwargs) 重寫 forward。setup_context 不被重寫。為 backward 設定 ctx 發生在 forward 內部。

2. 使用簽名 forward(*args, **kwargs) 重寫 forward 並重寫 setup_context。為 backward 設定 ctx 發生在 setup_context 內部（而非 forward 內部）

有關更多詳細資訊，請參閱 torch.autograd.Function.forward() 和 擴充套件 torch.autograd。

返回型別

任意型別

static vjp(ctx, *grad_outputs)

定義使用後向模式自動微分來區分操作的公式。

所有子類都必須重寫此函式。（定義此函式等同於定義 vjp 函式。）

它必須接受一個上下文 ctx 作為第一個引數，後跟與 forward() 返回的輸出一樣多的引數（對於 forward 函式的非 Tensor 輸出，將傳入 None），並且它應該返回與 forward() 輸入的 Tensor 一樣多的 Tensor。每個引數是相對於給定輸出的梯度，每個返回值應該是相對於相應輸入的梯度。如果輸入不是 Tensor 或者是不需要梯度的 Tensor，則可以簡單地為該輸入傳遞 None 作為梯度。

上下文可用於檢索在前向過程中儲存的 Tensor。它還有一個屬性 ctx.needs_input_grad，它是一個布林值元組，表示每個輸入是否需要梯度。例如，如果 forward() 的第一個輸入需要計算相對於輸出的梯度，則 backward() 將具有 ctx.needs_input_grad[0] = True。

返回型別

任意型別

static 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() 的 randomness 選項。

• 一個 in_dims 元組作為第二個引數。對於 args 中的每個 arg，in_dims 都有一個相應的 Optional[int]。如果 arg 不是 Tensor 或者 arg 未進行 vmap 操作，則為 None；否則，它是一個整數，指定 Tensor 的哪個維度正在進行 vmap 操作。

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

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

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