This PR adds a JAX-based CMA-ES version, which runs on different hardware backends (CPU, GPU, TPU). It overlaps with #6, but does not run into OOM issues, etc. I hope no one is mad at me for pushing this, since there was no update on the other PR for a month.
evosax
Source Code: https://github.com/RobertTLange/evosax/blob/main/evosax/strategies/cma_es.py