I love JAX. It can be a great replacement for Numba or whatever, the @jit works really well. And vmap is amazing... I often get lost in the matrix sauce when

I love JAX. It can be a great replacement for Numba or whatever, the @jit works really well. And vmap is amazing... I often get lost in the matrix sauce when batching, using JAX I develop as if it's just a single instance, then vmap that shit. The ecosystem I was using was: jax, optax, haiku. The big issue I had was: I was developing on the CPU, then moved to running it on a GPU, and it wasn't as fast as I expected-- I started debugging, and saw there was still lots of communication between the CPU and GPU even tho it was all jit'd. I think PyTorch is a more user friendly for writing high performance models if you're not straying too far from the beaten path. But I really love JAX would like to play around with it more to understand these pits I'm falling into. And another complaint is I can't run it on my Macbook M1 GPU... but I'm seeing this page now, so maybe that's not true anymore: https://developer.apple.com/metal/jax/ https://developer.apple.com/metal/jax/