jvp

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

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

source edit

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 useful for propagating tangent information (sensitivities) through a network, for directional derivatives, and as a building block for second-order methods.

Computed exactly with two reverse-mode passes (the double-vjp trick): $g(u) = J^\top u$ is linear in a dummy cotangent $u$ , so its vector-Jacobian product with $v$ is $(J^\top)^\top v = J v$ — or, through an op whose backward has no graph mode (LSTM, a padded avg_pool), one reverse pass per output element. No finite differences are involved.

Parameters

funccallable

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

inputsTensor or tuple of Tensor

Primal point

x

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

vTensor or tuple of Tensor= None

Tangent vector(s), one per input, each of its input's shape and device. May be left out when inputs is one single-element tensor, which is then seeded with 1.

create_graphbool= False

If True the returned outputs and tangents are differentiable with respect to the inputs that require grad (and v). Defaults to False.

strictbool= False

If True, refuse an output that does not require grad, an input the outputs do not depend on and an output whose tangent would be zeros for want of a path — and, under create_graph, a tangent that is a constant. Defaults to False.

Returns

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

(primals_out, tangents_out) where primals_out = func(*inputs) — detached unless create_graph — and tangents_out has the same structure and shapes as primals_out and holds $J v$ . An output that does not depend on the inputs has a zero tangent.

Raises

RuntimeError

v has the wrong number of entries, or is left out when inputs is not one single-element tensor, or strict refuses a result.

NotImplementedError

create_graph through an op whose backward has no graph mode.

(ShapeMismatch, DeviceMismatch)

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

Notes

The complementary operation is vjp, which computes $v^\top J$ with one reverse pass. This JVP costs two.

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)
>>> tangent
tensor([2., 0., 0.])

Used by 2

>>> 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) >>> tangent tensor([2., 0., 0.])