functorch.jacfwd

functorch.jacfwd(func, argnums=0, has_aux=False, *, randomness='error')

使用正向模式自動微分計算 func 關於索引 argnum 處的參數的雅可比矩陣

參數

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

• argnums（int 或 Tuple[int]）– 可選，整數或整數元組，指定要計算雅可比矩陣的參數。預設值：0。

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

• randomness（str）– 標記，指示要使用的隨機性類型。詳見 vmap()。允許的值："different"、"same"、"error"。預設值："error"

回傳值

返回一個函數，該函數接受與 func 相同的輸入，並返回 func 關於 argnums 指定參數的雅可比矩陣。如果 has_aux 為 True，則返回的函數會返回一個 (jacobian, aux) 元組，其中 jacobian 是雅可比矩陣，aux 是 func 返回的輔助物件。

注意事項

您可能會看到此 API 因為「運算子 X 未實作正向模式自動微分」而錯誤。如果是這樣，請提交錯誤報告，我們會優先處理。另一種方法是使用 jacrev()，它有更好的運算子覆蓋範圍。

對點態一元運算的基本用法將會得到一個對角矩陣作為雅可比矩陣

>>> from torch.func import jacfwd
>>> x = torch.randn(5)
>>> jacobian = jacfwd(torch.sin)(x)
>>> expected = torch.diag(torch.cos(x))
>>> assert torch.allclose(jacobian, expected)

jacfwd() 可以與 vmap 組合以產生批次的雅可比矩陣

>>> from torch.func import jacfwd, vmap
>>> x = torch.randn(64, 5)
>>> jacobian = vmap(jacfwd(torch.sin))(x)
>>> assert jacobian.shape == (64, 5, 5)

如果您想計算函數的輸出以及函數的雅可比矩陣，請使用 has_aux 標記將輸出作為輔助物件返回

>>> from torch.func import jacfwd
>>> x = torch.randn(5)
>>>
>>> def f(x):
>>>   return x.sin()
>>>
>>> def g(x):
>>>   result = f(x)
>>>   return result, result
>>>
>>> jacobian_f, f_x = jacfwd(g, has_aux=True)(x)
>>> assert torch.allclose(f_x, f(x))

此外，jacrev() 可以與自身或 jacrev() 組合以產生黑塞矩陣

>>> from torch.func import jacfwd, jacrev
>>> def f(x):
>>>   return x.sin().sum()
>>>
>>> x = torch.randn(5)
>>> hessian = jacfwd(jacrev(f))(x)
>>> assert torch.allclose(hessian, torch.diag(-x.sin()))

預設情況下，jacfwd() 計算關於第一個輸入的雅可比矩陣。然而，它可以透過使用 argnums 計算關於不同參數的雅可比矩陣

>>> from torch.func import jacfwd
>>> def f(x, y):
>>>   return x + y ** 2
>>>
>>> x, y = torch.randn(5), torch.randn(5)
>>> jacobian = jacfwd(f, argnums=1)(x, y)
>>> expected = torch.diag(2 * y)
>>> assert torch.allclose(jacobian, expected)

此外，將元組傳遞給 argnums 將計算關於多個參數的雅可比矩陣

>>> from torch.func import jacfwd
>>> def f(x, y):
>>>   return x + y ** 2
>>>
>>> x, y = torch.randn(5), torch.randn(5)
>>> jacobian = jacfwd(f, argnums=(0, 1))(x, y)
>>> expectedX = torch.diag(torch.ones_like(x))
>>> expectedY = torch.diag(2 * y)
>>> assert torch.allclose(jacobian[0], expectedX)
>>> assert torch.allclose(jacobian[1], expectedY)

警告

我們已將 functorch 整合到 PyTorch 中。作為整合的最後一步，functorch.jacfwd 從 PyTorch 2.0 開始已被棄用，並將在 PyTorch >= 2.3 的未來版本中刪除。請改用 torch.func.jacfwd；有關更多詳細資訊，請參閱 PyTorch 2.0 版本說明和/或 torch.func 遷移指南 https://pytorch.com.tw/docs/stable/func.migrating.html