Add a linearize like jax.linearize

[zou3519]

jvp(f, x, tangents) computes both f(x) and the Jacobian-vector product.

The goal of linearize(f) is to return a function that can compute forward-mode AD multiple times without re-evaluating f(x).

[TJHeeringa]

Is there a workaround until a function like this has been properly implemented?

[zou3519]

A workaround is to use jvp -- jvp will re-evaluate the function you pass to it every time, so the performance won't be as good.

[kshitij12345]

Added in https://github.com/pytorch/pytorch/pull/94173