grad_and_value

→Callable

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

source

Return a function computing both the gradient and the primal value.

Equivalent to combining grad with an evaluation of func, but evaluates the forward pass exactly once. This is the canonical pattern for training loops where the loss value is also needed for logging or learning-rate scheduling.

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) and the wrapped callable returns (grads, (loss, aux)). Default False.

Returns

Callable

Function returning (grads, value) — or (grads, (value, aux)) when has_aux=True.

Notes

Mathematically computes both $f(x)$ and $\nabla f(x)$ in a single forward + backward sweep, saving redundant work compared to calling func and grad(func) separately.

Examples

>>> import lucid
>>> from lucid.func import grad_and_value
>>> f = lambda x: (x ** 2).sum()
>>> gv = grad_and_value(f)
>>> grads, value = gv(lucid.tensor([1.0, 2.0, 3.0]))
>>> value  # 14.0
Tensor(14.)