Yes, you can just pass the function and it runs the optimization, in severe cases, you can start by doing a grid `.scan` or `.simplex` (Nelder-Mead simplex method), then `migrad` to minimize, and `.hesse` for 1 sigma bound.
You can also provide a gradient function that, well computes the gradient instead of computing it numerically.
> I imagine you can also do this in numba
Yes! It compiles to native code.
> Another big speedup in Jax is due to vmap/pmap, which allow to vectorise/parallelise computation.
It's possible to compile some vmapped functions with numba too, if the link is any indication, you may see even greater speedup than just jax.grad
However, I do concur that jax's vmap is absolutely fantastic and I found it very useful on many occasions.