I'm happy to hear about Equinox being used! (I'm the author.)
I'm curious what your workloads are that you're seeing speedups of as much as 1e4? Greatest I've heard of before was ~1e2 on some differential equation solving.
I'm curious what your workloads are that you're seeing speedups of as much as 1e4? Greatest I've heard of before was ~1e2 on some differential equation solving.
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_Optimi..., see page 87.