When to use `jax.jacrev()` vs `jax.jacfwd()`?

[arcruz0]

jax.jacrev() should be a performant option for calls with many parameters and few outputs (for example, a predictions(..., by = TRUE) call for a model with many parameters). How to handle which function to use?

[arcruz0]

^ Implemented jacx.jacrev() for predictions(, by = T), although the difference appears to be minimal with something like K=1000. It might kick-in with a larger K?

[arcruz0]

^ Implemented jacx.jacrev() for predictions(, by = var). Still little difference performance-wise.