Welcome to evosax

Basic evosax API Usage 🍲

import jax
from evosax import CMA_ES
from evosax.problems import batch_rosenbrock

# Instantiate the search strategy
rng = jax.random.PRNGKey(0)
strategy = CMA_ES(popsize=20, num_dims=2, elite_ratio=0.5)
params = strategy.default_params
state = strategy.initialize(rng, params)

# Run the ask-eval-tell loop
for t in range(num_generations):
    rng, rng_gen = jax.random.split(rng)
    x, state = strategy.ask(rng_gen, state, params)
    fitness = batch_rosenbrock(x, 1, 100)
    state = strategy.tell(x, fitness, state, params)

# Get best overall population member & its fitness
state["best_member"], state["best_fitness"]