evosax Release πŸŽ‰ - Evolution Strategies in JAX 🦎


I am more than excited to share that I have just released evosax – a JAX-based library of evolution strategies. evosax allows you to leverage JAX, XLA compilation and auto-vectorization/parallelization to scale ES to your favorite accelerators. The API is based on the classical ask, evaluate, tell cycle of ES. Both ask and tell calls are compatible with jit, vmap/pmap and lax.scan. It includes a vast set of both classic (e.g. CMA-ES, Differential Evolution, etc.) and modern neuroevolution (e.g. OpenAI-ES, Augmented RS, etc.) strategies. πŸ‘‰ Colab