| > I haven't tried numba but I've heard good things about it! Nice linked tutorial. If I understand correctly, you pass in jitted functions (using numba, and jax) into iminuit which does the optimisation? 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. |