functorch.jvp¶

functorch.jvp(func, primals, tangents, *, strict=False, has_aux=False)[source]¶

代表雅可比-向量积，返回一个包含func(*primals)输出和“在primals处评估的func的雅可比”乘以tangents的元组。这也被称为前向模式自动微分。

Parameters

func (function) – 一个Python函数，它接受一个或多个参数，其中一个必须是Tensor，并返回一个或多个Tensor
primals (张量) – 必须是func的位置参数，且必须全部为张量。返回的函数也将计算相对于这些参数的导数。
切线 (张量) – 用于计算雅可比-向量积的“向量”。必须与func的输入具有相同的结构和大小。
has_aux (bool) – 标志，表示 func 返回一个 (output, aux) 元组，其中第一个元素是要微分的函数的输出，第二个元素是不会被微分的其他辅助对象。默认值：False。

Returns

返回一个包含在primals处评估的func输出和雅可比向量积的(output, jvp_out)元组。如果has_aux is True，则返回一个(output, jvp_out, aux)元组。

注意

你可能会看到这个API错误提示“forward-mode AD not implemented for operator X”。如果是这样，请提交一个错误报告，我们会优先处理。

jvp 在你希望计算函数 R^1 -> R^N 的梯度时非常有用

>>> from torch.func import jvp
>>> x = torch.randn([])
>>> f = lambda x: x * torch.tensor([1., 2., 3])
>>> value, grad = jvp(f, (x,), (torch.tensor(1.),))
>>> assert torch.allclose(value, f(x))
>>> assert torch.allclose(grad, torch.tensor([1., 2, 3]))

jvp() 可以通过为每个输入传递切线来支持具有多个输入的函数

>>> from torch.func import jvp
>>> x = torch.randn(5)
>>> y = torch.randn(5)
>>> f = lambda x, y: (x * y)
>>> _, output = jvp(f, (x, y), (torch.ones(5), torch.ones(5)))
>>> assert torch.allclose(output, x + y)

警告

我们已经将functorch集成到PyTorch中。作为集成的最后一步，functorch.jvp在PyTorch 2.0中已被弃用，并将在未来版本PyTorch >= 2.3中删除。请改用torch.func.jvp；更多详情请参阅PyTorch 2.0发布说明和/或torch.func迁移指南https://pytorch.org/docs/master/func.migrating.html