Evolving Neural Networks in JAX

“So why should I switch from to JAX?”. The classic first passive-aggressive question when talking about the new ‘kid on the block’. Here is my answer: JAX is not simply a fast library for automatic differentiation. If your scientific computing project wants to benefit from XLA, JIT-compilation and the bulk-array programming paradigm – then JAX provides a wonderful API.