torch.func.jacfwd

torch.func.jacfwd(func, argnums=0, has_aux=False, *, randomness='error') [исходный код]

Вычисляет якобиан func относительно аргумента(ов) с индексом argnum с помощью автоматического дифференцирования в прямом режиме

Параметры:

• func (function) – Функция Python, которая принимает один или несколько аргументов, один из которых должен быть Tensor, и возвращает один или несколько Tensor
• argnums (int или tuple[int, ...]) – Необязательное целое число или кортеж целых чисел, указывающий, для каких аргументов нужно получить якобиан. По умолчанию: 0.
• has_aux (bool) – Флаг, указывающий, что func возвращает кортеж (output, aux), в котором первый элемент — результат функции, для которой вычисляется производная, а второй элемент — вспомогательные объекты, для которых производная вычисляться не будет. По умолчанию: False.
• randomness (str) – Флаг, указывающий, какой тип случайности использовать. Дополнительные сведения см. в разделе vmap(). Допустимые значения: “different”, “same”, “error”. По умолчанию: “error”

Возвращает:

Возвращает функцию, принимающую те же входные данные, что и func, и возвращающую якобиан func относительно аргумента(ов) с индексом argnums. Если has_aux is True, возвращаемая функция вместо этого возвращает кортеж (jacobian, aux), где jacobian — якобиан, а aux — вспомогательные объекты, возвращённые func.

Тип возвращаемого значения:

Callable[…, Any]

Примечание

Эта функция API может завершиться ошибкой «forward-mode AD not implemented for operator X» («прямая автоматическая дифференциация не реализована для оператора 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)