torch.func.jacfwd

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

使用前向模式自動微分計算 func 相對於 argnum 索引處的引數（們）的雅可比矩陣。

引數

• func (function) – A Python function that takes one or more arguments, one of which must be a Tensor, and returns one or more Tensors

• argnums (int or Tuple[int]) – 可選，整數或整數元組，用於指定要計算雅可比矩陣的引數。預設值：0。

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

• randomness (str) – 指示要使用的隨機性型別的標誌。更多詳細資訊請參閱 vmap()。允許的值：“different”、“same”、“error”。預設值：“error”。

返回

返回一個函式，該函式接受與 func 相同的輸入，並返回 func 相對於 argnums 中指定引數的 Jacobian。如果 has_aux 為 True，則返回的函式將返回一個 (jacobian, aux) 元組，其中 jacobian 是 Jacobian，aux 是 func 返回的輔助物件。

注意

您可能會看到此 API 因“前向模式 AD 未對運算子 X 實現”而報錯。如果遇到這種情況，請提交一個 Bug 報告，我們會優先處理。另一種選擇是使用 jacrev()，它的運算子覆蓋範圍更廣。

A basic usage with a pointwise, unary operation will give a diagonal array as the Jacobian

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

If you would like to compute the output of the function as well as the jacobian of the function, use the has_aux flag to return the output as an auxiliary object

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

Additionally, passing a tuple to argnums will compute the Jacobian with respect to multiple arguments

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