jvp
→tuplejvp(func: Callable[..., Tensor | tuple[Tensor, ...]], primals: tuple[Tensor, ...], tangents: tuple[Tensor, ...], strict: bool = False)Compute the Jacobian-vector product of func at primals.
Returns both the primal output func(*primals) and the directional
derivative — 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
funcCallableprimalstuple of Tensorfunc and its Jacobian.tangentstuple of Tensorstrictbool= FalseTrue, refuse what lucid.autograd.jvp refuses under
strict — an output independent of the inputs among them —
instead of giving it a zero tangent. Default False.Returns
tuple(primals_out, tangents_out) with the same nested structure as
func(*primals).
Raises
RuntimeErrorprimals and tangents differ in length, or strict
refuses a result.NotImplementedErrorNotes
is exact, not a finite difference: it is computed by
lucid.autograd.jvp in two reverse passes (the double-VJP
trick), or — through an op whose backward has no graph mode, such as
LSTM or a padded avg_pool — in one reverse pass per output
element.
The results join the caller's autograd graph as the reference's do.
Under no_grad neither does. Otherwise the primal output is attached
when it depends on a tensor that requires grad — a primal, or a
parameter func captures — and the tangent is attached when a
primal, a tangent or such a parameter requires grad. Through an op
with no graph-mode backward the tangent cannot be attached: when only
captured parameters would attach it, it comes back detached with a
warning.
func runs once when grad mode is off or a primal requires grad.
With grad mode on and no primal requiring grad it runs twice: once
at the primals, which is the only way to learn whether the output
reaches the caller's graph through tensors func captures, and once
for the tangent. The second run replays the first one's random draws,
so a dropout mask is the same in both, but any other side effect
happens twice — a BatchNorm in training mode updates its running
statistics twice, and an in-place buffer update is applied twice.
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
tensor([2., 4., 6.])