Use `jax.linearize` instead of `jax.jvp`
Closed this issue · 1 comments
mfschubert commented
It seems that using jax.linearize would be more efficient than jax.jvp in the sparsejac implementation of jacfwd. We should investigate this and change the implementation.
https://jax.readthedocs.io/en/latest/_autosummary/jax.linearize.html
mfschubert commented
Addressed by #13