vjp

→tuple of (Tensor or tuple of Tensor, Tensor or tuple of Tensor)

vjp(func: Callable[..., Tensor | tuple[Tensor, ...]], inputs: Tensor | tuple[Tensor, ...], v: Tensor | tuple[Tensor, ...] | None = None, create_graph: bool = False, strict: bool = False)

source edit

Vector-Jacobian product $v^\top J$ (reverse-mode AD).

Given $f : \mathbb{R}^n \to \mathbb{R}^m$ with Jacobian $J \in \mathbb{R}^{m \times n}$ and cotangent vector $v \in \mathbb{R}^m$ , returns $v^\top J \in \mathbb{R}^{n}$ along with the primal output $y = f(x)$ . Computing a full VJP costs one backward pass.

Parameters

funccallable

Function mapping Tensor inputs to a Tensor (or tuple thereof).

inputsTensor or tuple of Tensor

Primal point

x

at which

J

is evaluated. Left as it is: the product is taken at copies that require grad.

vTensor or tuple of Tensor= None

Cotangent vector(s), one per output of func, each of its output's shape and device — the seed rules of lucid.autograd.grad. May be left out when func returns one single-element tensor.

create_graphbool= False

If True the returned VJP is itself differentiable. Defaults to False.

strictbool= False

If True, refuse an output that does not require grad and an input the outputs do not depend on, instead of giving it a zero gradient. Defaults to False.

Returns

tuple of (Tensor or tuple of Tensor, Tensor or tuple of Tensor)

(output, vjp) where output = func(*inputs) — detached unless create_graph — and vjp is $v^\top J$ , shaped and packed like inputs: one tensor for one tensor, a tuple for a tuple. An input the outputs do not depend on gets zeros.

Raises

RuntimeError

The number of entries of v differs from the number of outputs, v is left out and func does not return one single-element tensor, or strict refuses a result.

(ShapeMismatch, DeviceMismatch)

An entry of v does not have its output's shape or device.

Notes

The dual to vjp is jvp, which computes $J v$ .

Examples

>>> import lucid
>>> from lucid.autograd import vjp
>>> x = lucid.tensor([1.0, 2.0, 3.0])
>>> v = lucid.tensor([1.0, 1.0, 1.0])
>>> y, grad_x = vjp(lambda x: x * x, x, v)
>>> grad_x
tensor([2., 4., 6.])

Used by 2

>>> import lucid >>> from lucid.autograd import vjp >>> x = lucid.tensor([1.0, 2.0, 3.0]) >>> v = lucid.tensor([1.0, 1.0, 1.0]) >>> y, grad_x = vjp(lambda x: x * x, x, v) >>> grad_x tensor([2., 4., 6.])