Jax is 10x faster than a few months ago!

[GeoffNN]

Hi all,

Want to say wow, and thanks! I re-ran code that I wrote in April this year, on GPU. It uses flax and jax's gmres solver. I didn't change anything on my side apart from updating jax, jaxopt, flax, optax and chex, and I'm getting 10x speedups. The same exact code used to take 5h to run, and now it's 30min! (!!!!)
Not sure where exactly the improvement comes from, but it's awesome.

So thanks 🥳

[mattjj]

Thanks so much for letting us know! We've been working hard on making things faster behind the scenes, and actually the best is yet to come. At NVIDIA GTC 2022 a couple weeks ago we announced a new collaboration between Google and NVIDIA to improve JAX further on GPUs, especially around performance.

[GeoffNN]

Can't wait to rerun the same code in a few months 😄 Awesome work!!

[tetterl]

Amazing work! I assume the improvements that are already in are mainly on GPU and not CPU?