  Evolutionary computation has been shown to be a highly effective method for
training neural networks, particularly when employed at scale on CPU clusters.
Recent work have also showcased their effectiveness on hardware accelerators,
such as GPUs, but so far such demonstrations are tailored for very specific
tasks, limiting applicability to other domains. We present EvoJAX, a scalable,
general purpose, hardware-accelerated neuroevolution toolkit. Building on top
of the JAX library, our 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.
We provide extensible examples of EvoJAX for a wide range of tasks, including
supervised learning, reinforcement learning and generative art. Since EvoJAX
can find solutions to most of these tasks within minutes on a single
accelerator, compared to hours or days when using CPUs, our toolkit can
significantly shorten the iteration cycle of evolutionary computation
experiments. EvoJAX is available at https://github.com/google/evojax
