linearize

→tuple

linearize(func: Callable[..., Tensor | tuple[Tensor, ...]], primals: Tensor = ())

source edit

Linearise func around primals for reuse across many tangents.

Returns the primal output together with a callable linear_fn that applies the first-order Taylor expansion of func at primals. Mathematically linear_fn(t) evaluates $J_f(\mathrm{primals}) \, t$ — the same quantity as jvp — but the linearisation cost is paid once even if many tangent vectors are queried subsequently.

Parameters

funcCallable

Differentiable function.

*primalsTensor= ()

Points at which to linearise func.

Returns

tuple

(primals_out, linear_fn). Calling linear_fn(*tangents) returns the JVP at primals against tangents.

Notes

Conceptually equivalent to a Taylor expansion truncated at first order:

f(x + t) \approx f(x) + J_f(x) \, t.

func runs once here and once per call of linear_fn. Whether the results join the caller's autograd graph follows jvp, and is decided here, from the primal output, once. The returned linear_fn is flagged for vmap isolation so that composing vmap(linear_fn) slices tangents one at a time.

Examples

>>> import lucid
>>> from lucid.func import linearize
>>> f = lambda x: x ** 2
>>> x = lucid.tensor([1.0, 2.0, 3.0])
>>> y, lin = linearize(f, x)
>>> lin(lucid.ones_like(x))  # 2 * x
tensor([2., 4., 6.])

>>> import lucid >>> from lucid.func import linearize >>> f = lambda x: x ** 2 >>> x = lucid.tensor([1.0, 2.0, 3.0]) >>> y, lin = linearize(f, x) >>> lin(lucid.ones_like(x)) # 2 * x tensor([2., 4., 6.])