2 ms·

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

• chillee 5y ago
  If you’re massively dominated by overhead I can see it. I’ve definitely done microbenchmarks before where compilation gets you 1000x improvement.
• PartiallyTyped 5y ago
  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.