jacfwd

→Callable

jacfwd(func: Callable[..., Tensor | tuple[Tensor, ...]], argnums: int | tuple[int, ...] = 0, has_aux: bool = False, randomness: str = 'error')

source

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

Materialises the Jacobian one column at a time by repeatedly calling jvp with one-hot tangent vectors. Each column corresponds to the partial derivative of every output with respect to a single input coordinate.

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.

randomnessstr= 'error'

Randomness policy forwarded to internal vmap calls. Default "error".

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 input dimension $n$ (one forward / JVP per column). Prefer this transform over jacrev when $n \ll m$ .

Examples

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