jacobian

→Tensor or tuple of Tensor

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

source

Compute the Jacobian matrix of func with respect to each input.

The Jacobian of a vector-valued function $f : \mathbb{R}^n \to \mathbb{R}^m$ is

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

Lucid evaluates it row-by-row by repeated reverse-mode backward passes — one per output element — seeding each pass with a one-hot cotangent so the resulting input gradient is exactly the corresponding Jacobian row. The cost therefore scales with the output dimension $m$ ; prefer vjp when only $v^\top J$ is needed and jvp when only $J v$ is needed.

Parameters

funccallable

A function mapping Tensor inputs to a Tensor (or tuple of Tensor). Must be differentiable w.r.t. each positional input.

inputsTensor or tuple of Tensor

Input tensor(s) at which the Jacobian is evaluated. They are silently promoted to requires_grad=True if needed.

create_graphbool= False

If True the Jacobian itself is differentiable, enabling higher-order derivatives (e.g. building hessian on top). Defaults to False.

strictbool= False

Reserved for stricter shape/dtype validation. Currently unused.

vectorizebool= False

Reserved for a future vmap-based implementation. Currently unused.

Returns

Tensor or tuple of Tensor

For a single input x the returned tensor has shape (prod(out_shape), prod(x.shape)). For multiple inputs a tuple is returned, one Jacobian block per input.

Notes

Reverse-mode differentiation makes the cost per row $O(\text{cost}(f))$ ; the full Jacobian therefore costs $O(m \cdot \text{cost}(f))$ . For square or wide Jacobians ( $m \ge n$ ) forward-mode would be cheaper — Lucid does not yet ship a forward-mode implementation, so this routine is preferred for tall Jacobians ( $m \ll n$ ).

Examples

>>> import lucid
>>> from lucid.autograd import jacobian
>>> x = lucid.tensor([1.0, 2.0, 3.0])
>>> def f(x):
...     return x * x
>>> J = jacobian(f, x)
>>> J.shape
(3, 3)