jvp

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

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

source

Jacobian-vector product $J v$ (forward-mode directional derivative).

Given $f : \mathbb{R}^n \to \mathbb{R}^m$ with Jacobian $J \in \mathbb{R}^{m \times n}$ and a tangent vector $v \in \mathbb{R}^n$ , returns

J v = \left.\frac{d}{dt} f(x + t v)\right|_{t=0} \in \mathbb{R}^{m}

along with the primal output $y = f(x)$ . JVPs are the natural primitive of forward-mode AD and are useful for propagating tangent information (sensitivities) through a network in a single forward sweep, for computing directional derivatives, and as a building block for second-order methods.

Lucid currently realises the JVP via a symmetric central finite difference

J v \approx \frac{f(x + \varepsilon v) - f(x - \varepsilon v)} {2 \varepsilon},

with $\varepsilon = 10^{-4}$ . This avoids the need for a true forward-mode implementation while still being accurate enough for testing and most applications.

Parameters

funccallable

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

inputsTensor or tuple of Tensor

Primal point

x

vTensor or tuple of Tensor

Tangent vector(s) matching the input shape(s).

create_graphbool= False

Reserved for the future native forward-mode implementation. Currently unused.

strictbool= False

Reserved for stricter validation. Currently unused.

Returns

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

(primals_out, tangents_out) where primals_out = func(*inputs) and tangents_out has the same shape as primals_out and holds $J v$ .

Notes

The complementary operation is vjp, which computes $v^\top J$ cheaply via reverse-mode. Use jvp when the input dimension is small relative to the output dimension; otherwise reverse-mode is more efficient.

Examples

>>> import lucid
>>> from lucid.autograd import jvp
>>> x = lucid.tensor([1.0, 2.0, 3.0])
>>> v = lucid.tensor([1.0, 0.0, 0.0])
>>> def f(x):
...     return x * x
>>> y, tangent = jvp(f, x, v)