torch.autograd.functional.jvp

torch.autograd.functional.jvp(func, inputs, v=None, create_graph=False, strict=False) [исходный код]

Вычисляет скалярное произведение между якобианом заданной функции в точке, заданной входными данными, и вектором v.

Параметры:

• func (function) – функция Python, которая принимает входные данные Tensor и возвращает кортеж Tensor или Tensor.
• inputs (tuple of Tensors or Tensor) – входные данные функции func.
• v (tuple of Tensors or Tensor) – Вектор, для которого вычисляется произведение якобиана на вектор. Должен иметь тот же размер, что и входные данные func. Этот аргумент необязателен, если входные данные func содержат один элемент; если аргумент не указан, он будет задан как Tensor, содержащий одно значение 1.
• create_graph (bool, optional) – Если True, и выходные данные, и результат будут вычислены дифференцируемым способом. Обратите внимание: если strict имеет значение False, результат не может требовать вычисления градиентов или быть не связанным с входными данными. Значение по умолчанию — False.
• strict (bool, optional) – Если True, будет вызвана ошибка, если мы обнаружим входные данные, от которых не зависят все выходные данные. Если False, для таких входных данных мы вернём Tensor нулей в качестве jvp — это ожидаемое математическое значение. Значение по умолчанию — False.

Возвращает:

кортеж, содержащий:

func_output (кортеж Tensor или Tensor): выходные данные func(inputs)

jvp (кортеж Tensor или Tensor): результат скалярного произведения с той же формой, что и выходные данные.

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

output (tuple)

Примечание

autograd.functional.jvp вычисляет jvp, используя обратное распространение для обратного прохода (иногда это называют трюком двойного обратного прохода). Это не самый производительный способ вычисления jvp. Вместо этого рассмотрите возможность использования torch.func.jvp() или низкоуровневого API прямого автоматического дифференцирования.

Пример

>>> def exp_reducer(x):
...     return x.exp().sum(dim=1)
>>> inputs = torch.rand(4, 4)
>>> v = torch.ones(4, 4)
>>> jvp(exp_reducer, inputs, v)
(tensor([6.3090, 4.6742, 7.9114, 8.2106]),
 tensor([6.3090, 4.6742, 7.9114, 8.2106]))

>>> jvp(exp_reducer, inputs, v, create_graph=True)
(tensor([6.3090, 4.6742, 7.9114, 8.2106], grad_fn=<SumBackward1>),
 tensor([6.3090, 4.6742, 7.9114, 8.2106], grad_fn=<SqueezeBackward1>))

>>> def adder(x, y):
...     return 2 * x + 3 * y
>>> inputs = (torch.rand(2), torch.rand(2))
>>> v = (torch.ones(2), torch.ones(2))
>>> jvp(adder, inputs, v)
(tensor([2.2399, 2.5005]),
 tensor([5., 5.]))