vjp

→tuple

vjp(func: Callable[..., Tensor | tuple[Tensor, ...]], primals: Tensor = (), has_aux: bool = False)

source

Compute the vector-Jacobian product of func at primals.

Evaluates func at the supplied primal inputs and returns both the output and a callable vjp_fn that, given cotangent vectors $v$ , returns $v^\top J(\mathrm{primals})$ — the backward-mode contraction of the Jacobian against v. This is the workhorse used internally by grad and jacrev, and is useful directly when many backward passes are needed against the same forward computation.

Parameters

funcCallable

Differentiable function taking the primals as positional inputs.

*primalsTensor= ()

Points at which to linearise func.

has_auxbool= False

If True, func must return (output, aux); the call then yields ((output, aux), vjp_fn). Default False.

Returns

tuple

(outputs, vjp_fn). vjp_fn(*cotangents) returns a tuple of input-gradient tensors of the same shapes as primals.

Notes

For $f : \mathbb{R}^n \to \mathbb{R}^m$ and cotangent $v \in \mathbb{R}^m$ ,

\mathrm{vjp\_fn}(v) = v^\top J_f(x), \quad J_f(x) \in \mathbb{R}^{m \times n}.

Cost is one backward pass per call to vjp_fn — the forward graph is retained so multiple cotangents can be applied cheaply.

Examples

>>> import lucid
>>> from lucid.func import vjp
>>> f = lambda x: x ** 2
>>> x = lucid.tensor([1.0, 2.0, 3.0])
>>> y, vjp_fn = vjp(f, x)
>>> (grads,) = vjp_fn(lucid.ones_like(y))  # 2 * x