tensorplay.func.linearize

tensorplay.func.linearize(func: Callable, *primals) → tuple[Any, Callable]

Evaluates func at primals and returns its linear approximation there.

The primal evaluation happens once, here; the returned jvp_fn reuses that linearization, so applying it to many tangents costs only the tangent passes. Use it in place of repeated jvp() calls at a fixed point.

Returns:

(output, jvp_fn). jvp_fn takes tangents matching the structure of primals and returns the directional derivative of the output.

Example

>>> x = tensorplay.randn(5)
>>> out, jvp_fn = linearize(tensorplay.sin, x)
>>> tangent_out = jvp_fn(tensorplay.ones(5))