functorch.jvp

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

`jvp` 代表 Jacobian 向量積，返回一個元組，包含 `func(*primals)` 的輸出以及「在 `primals` 處評估的 `func` 的 Jacobian」乘以 `tangents` 的結果。這也稱為前向模式自動微分。

參數

• func (函式) – 一個 Python 函式，它接受一個或多個參數（其中一個必須是張量），並返回一個或多個張量。

• primals (張量) – 傳遞給 `func` 的位置參數，這些參數必須都是張量。返回的函式也將計算關於這些參數的導數。

• tangents (張量) – 計算 Jacobian 向量積所用的「向量」。必須與 `func` 的輸入具有相同的結構和大小。

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

返回值

返回一個 `(output, jvp_out)` 元組，包含在 `primals` 處評估的 `func` 的輸出以及 Jacobian 向量積。如果 `has_aux` 為 `True`，則返回一個 `(output, jvp_out, aux)` 元組。

注意事項

您可能會看到此 API 因為「運算子 X 未實作前向模式自動微分」而錯誤。如果是這樣，請提交錯誤報告，我們會優先處理。

當您希望計算函式 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]))

藉由傳入每個輸入的切向量，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.dev.org.tw/docs/master/func.migrating.html