grad

→Callable

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

source

Build a function returning the gradient of func.

The cornerstone of functional-style autograd: rather than calling .backward() and reading .grad, grad(func) produces a new callable that, when invoked, returns the gradient tensor directly. Transforms compose, so grad(grad(func)) yields a second derivative and vmap(grad(func)) computes per-sample gradients.

Parameters

funcCallable

Function returning a scalar Tensor (or (scalar, aux) when has_aux=True).

argnumsint or tuple of int= 0

Positional argument index/indices to differentiate. Default 0.

has_auxbool= False

If True, func must return (loss, aux); the wrapped callable then returns (grads, aux) with aux forwarded through without differentiation. Default False.

Returns

Callable

Function with the same signature as func returning the gradient tensor — or a tuple of gradients when argnums is a tuple.

Notes

For $f : \mathbb{R}^n \to \mathbb{R}$ , grad(f) realises

\nabla f : \mathbb{R}^n \to \mathbb{R}^n, \quad (\nabla f)_i (x) = \frac{\partial f}{\partial x_i}(x).

Implementation is reverse-mode AD: one forward + one backward pass through func, independent of input dimensionality.

Examples

>>> import lucid
>>> from lucid.func import grad
>>> f = lambda x: (x ** 3).sum()
>>> df = grad(f)
>>> df(lucid.tensor([1.0, 2.0, 3.0]))  # 3 * x ** 2
Tensor([ 3., 12., 27.])