NestedIOFunction

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

此类仅出于向后兼容性原因而存在。在任何新用例中，请使用 Function 而不是此类。

backward(*gradients)

共享的 backward 工具。

返回类型

任何

backward_extended(*grad_output)

用户定义的 backward。

forward(*args)

共享的 forward 工具。

返回类型

任何

forward_extended(*input)

用户定义的 forward。

static jvp(ctx, *grad_inputs)

定义使用前向模式自动微分来区分操作的公式。

此函数应由所有子类重写。它必须接受一个上下文 ctx 作为第一个参数，然后是 forward() 接收到的输入数量（对于 forward 函数的非张量输入将传入 None），并且它应该返回与 forward() 的输出数量相同的张量。每个参数是关于给定输入的梯度，每个返回值应该是关于相应输出的梯度。如果输出不是张量或函数对该输出不可微，则只需为该输入传递 None 作为梯度。

You can use the ctx object to pass any value from the forward to this functions.

返回类型

任何

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)

Save given tensors for a future call to jvp().

save_for_forward 最多应调用一次，在 setup_context() 或 forward() 方法中，并且所有参数都应该是张量。

在 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)

Set whether to materialize grad tensors. Default is True.

这应仅从 setup_context() 或 forward() 方法中调用。

如果为 True，未定义的 grad 张量将在调用 backward() 和 jvp() 方法之前扩展为全零张量。

示例

>>> 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 的前向传递。

Either

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 函数的非张量输出将传入 None），并且它应该返回与 forward() 的输入数量相同的张量。每个参数是关于给定输出的梯度，每个返回值应该是关于相应输入的梯度。如果输入不是张量或是一个不需要梯度的张量，则可以为该输入传递 None 作为梯度。

上下文可用于检索在 forward 过程中保存的张量。它还有一个 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() 的随机性选项。

• 第二个参数是一个 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。