jacobian(func: Callable[..., Tensor | tuple[Tensor, ...]], inputs: Tensor | tuple[Tensor, ...], create_graph: bool = False, strict: bool = False, vectorize: bool = False)Compute the Jacobian matrix of func with respect to each input.
The Jacobian of a vector-valued function is
Lucid evaluates it row by row with reverse-mode passes — one per
output element — seeding each pass with a row of the identity so the
resulting input gradient is exactly the corresponding Jacobian row.
The cost therefore scales with the output dimension ; prefer
vjp when only is needed and jvp when
only is needed.
Parameters
funccallableTensor inputs to a Tensor (or
tuple of Tensor). Must be differentiable w.r.t. each
positional input.requires_grad and .grad do not change.create_graphbool= FalseTrue the Jacobian itself is differentiable — with respect
to the inputs that require grad — enabling higher-order
derivatives. Defaults to False.strictbool= FalseTrue, refuse an output that does not require grad, or an
input some output does not depend on, instead of giving it a
block of zeros — and, under create_graph, a block that is a
constant. Defaults to False.vectorizebool= FalseReturns
Tensor, tuple of Tensor or tuple of tuple of TensorOne block per output out and input x, of shape
out.shape + x.shape. Packed over the inputs as inputs
was, and over the outputs as func returned them: a bare
tensor for one output and one input, J[i] when only one of
the two is a tuple, J[i][j] for output i and input j
when both are.
Notes
Reverse-mode differentiation makes the cost per row ; the full Jacobian therefore costs . An input the output does not depend on gets a block of zeros.
Examples
>>> import lucid
>>> from lucid.autograd import jacobian
>>> x = lucid.tensor([1.0, 2.0, 3.0])
>>> def f(x):
... return x * x
>>> jacobian(f, x)
tensor([[2., 0., 0.], [0., 4., 0.], [0., 0., 6.]])