jacrev

→Callable

jacrev(func: Callable[..., Tensor | tuple[Tensor, ...]], argnums: int | tuple[int, ...] = 0, has_aux: bool = False, chunk_size: int | None = None)

source

Build a function returning the reverse-mode Jacobian of func.

Assembles the full Jacobian matrix one row at a time, each row obtained from a backward pass seeded with a one-hot cotangent. For a scalar-output func this collapses to grad.

Parameters

funcCallable

Differentiable function returning a Tensor (or (output, aux) when has_aux=True).

argnumsint or tuple of int= 0

Argument(s) to differentiate with respect to. Default 0.

has_auxbool= False

Whether func returns auxiliary data. Default False.

chunk_sizeint= None

Reserved for future vmap-batched row computation.

Returns

Callable

Function returning the Jacobian tensor (or tuple of tensors) with shape output.shape + arg.shape.

Notes

For $f : \mathbb{R}^n \to \mathbb{R}^m$ produces

J_{ij} = \frac{\partial f_i}{\partial x_j}, \quad J \in \mathbb{R}^{m \times n}.

Cost scales with the output dimension $m$ (one backward pass per row). Prefer jacfwd when $n < m$ .

Examples

>>> import lucid
>>> from lucid.func import jacrev
>>> f = lambda x: lucid.stack([x.sum(), (x ** 2).sum()])
>>> jacrev(f)(lucid.tensor([1.0, 2.0, 3.0]))  # shape (2, 3)