mutax 0.2.0

Evolutionary optimization algorithms in JAX

Mutax

SciPy-like differential evolution for JAX

Fully jax.jitted optimization of any JAX-compatible function. Serial and parallel execution on CPU, GPU, and TPU.

Installation

pip install mutax

Quick start

import jax
import jax.numpy as jnp
from mutax import differential_evolution

def cost_function(x):
    return jnp.sum(x**2)

bounds = [(-5, 5)] * 10  # 10-dimensional problem with bounds for each dimension

result = differential_evolution(cost_function, bounds, key=jax.random.key(0))
print("Best solution:", result.x)
print("Objective value:", result.fun)

Documentation

The documentation is available at Read the Docs.