torch.func.jvp

torch.func.jvp(func, primals, tangents, *, strict=False, has_aux=False)

代表 Jacobian-向量積 (Jacobian-vector product)，會回傳一個 tuple，其中包含 func(*primals) 的輸出，以及 “在 primals 處評估的 func 的 Jacobian” 乘以 tangents 的結果。 這也稱為前向模式自動微分 (forward-mode autodiff)。

參數

• func (function) – 一個 Python 函式，接受一個或多個參數，其中一個必須是 Tensor，並回傳一個或多個 Tensor

• primals (Tensors) – 傳遞給 func 的位置參數，這些參數都必須是 Tensor。 回傳的函式也會計算關於這些參數的導數。

• tangents (Tensors) – 計算 Jacobian-向量積的「向量」。 必須與輸入到 func 的結構和大小相同。

• has_aux (bool) – 指示 func 回傳一個 (output, aux) tuple 的旗標，其中第一個元素是要微分的函式的輸出，第二個元素是不會被微分的其他輔助物件。 預設值：False。

回傳

回傳一個 (output, jvp_out) tuple，其中包含在 primals 處評估的 func 的輸出以及 Jacobian-向量積。 如果 has_aux is True，則改為回傳一個 (output, jvp_out, aux) tuple。

注意

您可能會看到此 API 發生錯誤，並顯示「operator X 未實作前向模式 AD」。 如果是這樣，請提交錯誤報告，我們會優先處理它。

當您希望計算函數 R^1 -> R^N 的梯度時，jvp 非常有用

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

透過傳入每個輸入的 tangent，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)