4 ms·

What about JAX?

• kavalg 5y ago
  JAX is really cool, but still somewhat immature. I would love to see it taking more ground and improving wrt e.g. integration with tensorboard and getting all the goodies we have in tensorflow. If you are looking for a higher level framework, I would recommend elegy [0] which is very close to the keras API. [0] https://github.com/poets-ai/elegy https://github.com/poets-ai/elegy
• PartiallyTyped 5y ago
  Jax is great, but there are some rough edges. I am using Jax for differentiable programming, and in many cases, I saw enormous speedup after jit, sometimes in the ballpark of 1e4. For Neural Networks, I use Equinox, and/or Elegy.
  • patrickkidger 5y ago
    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.
    • 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.