torch.autograd.functional.vjp — PyTorch 2.7 documentation (original) (raw)

torch.autograd.functional.vjp(func, inputs, v=None, create_graph=False, strict=False)[source][source]¶

Compute the dot product between a vector v and the Jacobian of the given function at the point given by the inputs.

Parameters

• func (function) – a Python function that takes Tensor inputs and returns a tuple of Tensors or a Tensor.
• inputs (tuple of Tensors or Tensor) – inputs to the function func.
• v (tuple of Tensors or Tensor) – The vector for which the vector Jacobian product is computed. Must be the same size as the output of func. This argument is optional when the output of funccontains a single element and (if it is not provided) will be set as a Tensor containing a single 1.
• create_graph (bool, optional) – If True, both the output and result will be computed in a differentiable way. Note that when strictis False, the result can not require gradients or be disconnected from the inputs. Defaults to False.
• strict (bool, optional) – If True, an error will be raised when we detect that there exists an input such that all the outputs are independent of it. If False, we return a Tensor of zeros as the vjp for said inputs, which is the expected mathematical value. Defaults to False.

Returns

tuple with:

func_output (tuple of Tensors or Tensor): output of func(inputs)

vjp (tuple of Tensors or Tensor): result of the dot product with the same shape as the inputs.

Return type

output (tuple)

Example

def exp_reducer(x): ... return x.exp().sum(dim=1) inputs = torch.rand(4, 4) v = torch.ones(4) vjp(exp_reducer, inputs, v) (tensor([5.7817, 7.2458, 5.7830, 6.7782]), tensor([[1.4458, 1.3962, 1.3042, 1.6354], [2.1288, 1.0652, 1.5483, 2.5035], [2.2046, 1.1292, 1.1432, 1.3059], [1.3225, 1.6652, 1.7753, 2.0152]]))

vjp(exp_reducer, inputs, v, create_graph=True) (tensor([5.7817, 7.2458, 5.7830, 6.7782], grad_fn=), tensor([[1.4458, 1.3962, 1.3042, 1.6354], [2.1288, 1.0652, 1.5483, 2.5035], [2.2046, 1.1292, 1.1432, 1.3059], [1.3225, 1.6652, 1.7753, 2.0152]], grad_fn=))

def adder(x, y): ... return 2 * x + 3 * y inputs = (torch.rand(2), torch.rand(2)) v = torch.ones(2) vjp(adder, inputs, v) (tensor([2.4225, 2.3340]), (tensor([2., 2.]), tensor([3., 3.])))