jvp

→tuple

jvp(func: Callable[..., Tensor | tuple[Tensor, ...]], primals: tuple[Tensor, ...], tangents: tuple[Tensor, ...], strict: bool = False)

source

Compute the Jacobian-vector product of func at primals.

Returns both the primal output func(*primals) and the directional derivative $J(\mathrm{primals}) \cdot \mathrm{tangents}$ — the forward-mode contraction of the Jacobian against a tangent vector. Forward-mode is preferred over reverse-mode when the input dimension is small relative to the output dimension.

Parameters

funcCallable

Differentiable function returning a Tensor or tuple of Tensors.

primalstuple of Tensor

Points at which to evaluate func and its Jacobian.

tangentstuple of Tensor

Tangent vectors, each matching the shape of the corresponding primal.

strictbool= False

If True, raise when an output is independent of any input. Default False.

Returns

tuple

(primals_out, tangents_out) with the same nested structure as func(*primals).

Notes

Implemented via the exact $\alpha$ -perturbation trick: introduce a scalar $\alpha$ with requires_grad=True, substitute x + alpha * t for each primal, and read $\partial \text{out} / \partial \alpha$ at $\alpha = 0$ from a single backward pass. No finite differences are used, so the result is exact up to floating-point rounding.

Examples

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