EvoJAX is a scalable, general purpose, hardware-accelerated neuroevolution toolkit. Built on top of the JAX library, this toolkit enables neuroevolution algorithms to work with neural networks running in parallel across multiple TPU/GPUs. EvoJAX achieves very high performance by implementing the evolution algorithm, neural network and task all in NumPy, which is compiled just-in-time to run on accelerators. This repo also includes several extensible examples of EvoJAX for a wide range of tasks, including supervised learning, reinforcement learning and generative art, demonstrating how EvoJAX can run your evolution experiments within minutes on a single accelerator, compared to hours or days when using CPUs. EvoJAX is implemented in JAX which needs to be installed first.

Evolutionary Computation is a computational intelligence technique inspired by natural evolution. This method begins with the development of a group of people that respond to an issue; then, evaluate and modify the possible set of solutions to accomplish the best available solution. To a significant extent, Evolutionary Computation has been providing a highly effective method for neural networks; it becomes more visible when it comes to deploying at scale on CPU clusters. In this article, I like to write about EvoJAX, a scalable, general-purpose, hardware-accelerated neuroevolution toolkit. EvoJAX is based on the JAX library and allows neuroevolution algorithms to perform with neural networks parallelly across multiple TPU/GPUs.