Mainly optimizers/solvers for multivariate functions. The 1e4 speedup was on a Trust Region optimizer. The algorithm was implemented to solve "the hard case"[1

Mainly optimizers/solvers for multivariate functions. The 1e4 speedup was on a Trust Region optimizer. The algorithm was implemented to solve "the hard case"[1] involves multiple Cholesky factorizations, a matrix inversion, an eigenvalue decomposition on each step, and a call to scipy.linalg.solve_triangular. Part of the speedup is likely from caching/avoiding recomputations things. Granted, I had to rewrite a lot of the code to accommodate jax's peculiarities around python semantics, and made extensive use of jax.lax.{fori_loop, while_loop, scan, cond}. [1] http://www.apmath.spbu.ru/cnsa/pdf/monograf/Numerical_Optimization2006.pdf http://www.apmath.spbu.ru/cnsa/pdf/monograf/Numerical_Optimi..., see page 87.