Skip to content

yippiez/tinygp

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

35 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

tinygp

tinygp is a genetic programming library designed around using tinygrads as a mean to represent programms particularly as UOP trees. It uses ask & tell API from (Collette et al., 2013). UOP trees can be evolved given strategies on user defined metrics. This allows easier simplification and efficent computation of each program.

Basic Usage

import numpy as np
from tinygp.evaluate import eval_uop
from tinygp.strategies import CMA_ES


# Instantiate the search strategy
strategy = CMA_ES(
    population_size=32,
    to_k=8,
    mutation_rate=0.25,
    crossover_rate=0.8,
    max_depth=4,
    seed=0,
    maximize=True,
)

# Target data (example: fit y = x^3 + x^2 + x)
x = np.linspace(-1.0, 1.0, 128)
y = x**3 + x**2 + x

# Initialize state (pass None to ask on first generation)
state = None
num_generations = 100

# Ask-Eval-Tell loop
for i in range(num_generations):
    # Generate a set of candidate programs
    population, state = strategy.ask(state)

    # Evaluate the fitness of the population
    preds = np.stack([eval_uop(program, x) for program in population], axis=0)
    mse = np.mean((preds - y) ** 2, axis=1)
    fitness = -mse  # maximize fitness == minimize MSE

    # Update the evolution strategy
    state = strategy.tell(state, fitness)

# Get best program
assert state is not None, "state must be initialized after at least one generation"
assert state.best_program is not None, "best_program must exist after fitness evaluation"
assert state.best_fitness is not None, "best_fitness must exist after fitness evaluation"
state.best_program, state.best_fitness

Installation

The current recommended way to install tinygp is from source.

Direct

# Python3 installation
python3 -m pip install git+https://github.com/yippiez/tinygp.git

# For uv
uv pip install git+https://github.com/yippiez/tinygp.git

From source

git clone https://github.com/yippiez/tinygp.git
cd tinygp

# uv installation 
uv sync
# or install just the package into the active environment
uv pip install -e .

Benchmarks

tinygp includes multiple strategy backends and can be tuned for different quality/speed tradeoffs.

Quick snapshot from the current benchmark run (14 targets: nguyen_1..8, koza_1..3, keijzer_1..3, generations=5), currently compared against gplearn:

Library / Strategy Avg Test MSE (lower is better) Speed vs gplearn (higher is better)
gplearn 0.023048 1.00x
tinygp[ASEBO] 0.022195 5.87x
tinygp[Sep_CMA_ES] 0.026559 6.95x

The accuracy and speed varias with different configurations. Best way to find optimal strategy is to just test multiple strategies with different parameters for a given problem.

For full benchmark tables, methodology, and simplify-mode comparisons, see BENCHMARK.md.

Run all benchmarks with uv run tinygp-benchmark --strategy all --target all --generations 5. Pick a single strategy with uv run tinygp-benchmark --strategy CMA_ES --target nguyen_1 --generations 5.

Examples

  • examples/automatic_kernel_synthesis.py shows GP-based kernel optimization starting from a tinygrad-defined reference kernel.

Implemented Evolution Strategies

Strategy Reference Code
Simple Evolution Strategy Rechenberg (1978) SimpleES
OpenAI-ES Salimans et al. (2017) Open_ES
CMA-ES Hansen & Ostermeier (2001) CMA_ES
Sep-CMA-ES Ros & Hansen (2008) Sep_CMA_ES
xNES Wierstra et al. (2014) xNES
SNES Wierstra et al. (2014) SNES
MA-ES Bayer & Sendhoff (2017) MA_ES
LM-MA-ES Loshchilov et al. (2017) LM_MA_ES
Rm_ES Li & Zhang (2017) Rm_ES
PGPE Sehnke et al. (2010) PGPE
ARS Mania et al. (2018) ARS
ESMC Merchant et al. (2021) ESMC
Persistent ES Vicol et al. (2021) PersistentES
Noise-Reuse ES Li et al. (2023) NoiseReuseES
CR-FM-NES Nomura & Ono (2022) CR_FM_NES
Guided ES Maheswaranathan et al. (2018) GuidedES
ASEBO Choromanski et al. (2019) ASEBO
Discovered ES Lange et al. (2023a) DiscoveredES
LES Lange et al. (2023a) LearnedES
EvoTF Lange et al. (2024) EvoTF_ES
iAMaLGaM-Full Bosman et al. (2013) iAMaLGaM_Full
iAMaLGaM-Univariate Bosman et al. (2013) iAMaLGaM_Univariate
Gradientless Descent Golovin et al. (2019) GradientlessDescent
Simulated Annealing Rasdi Rere et al. (2015) SimAnneal
Hill Climbing Rasdi Rere et al. (2015) HillClimbing
Random Search Bergstra & Bengio (2012) RandomSearch
SV-CMA-ES Braun et al. (2024) SV_CMA_ES
SV-OpenAI-ES Liu et al. (2017) SV_OpenES
Simple Genetic Algorithm Such et al. (2017) SimpleGA
MR15-GA Rechenberg (1978) MR15_GA
SAMR-GA Clune et al. (2008) SAMR_GA
GESMR-GA Kumar et al. (2022) GESMR_GA
LGA Lange et al. (2023b) LearnedGA
Diffusion Evolution Zhang et al. (2024) DiffusionEvolution
Differential Evolution Storn & Price (1997) DifferentialEvolution
Particle Swarm Optimization Kennedy & Eberhart (1995) PSO

Supported UOPs

tinygp evaluates a focused UOP subset to keep GP search stable, comparable, and efficient.

From tinygrad UOps, tinygp currently supports ADD, CONST, DEFINE_VAR, EXP2, FDIV, LOG2, MAX, MUL, NEG, POW, RECIPROCAL, SIN, SQRT, SUB, and TRUNC; it does not support AFTER, ALLREDUCE, AND, ASSIGN, BARRIER, BINARY, BIND, BITCAST, BUFFER, BUFFERIZE, BUFFER_VIEW, CALL, CAST, CAT, CMPEQ, CMPLT, CMPNE, CONTIGUOUS, CONTIGUOUS_BACKWARD, CONTRACT, COPY, CUSTOM, CUSTOMI, DEFINE_LOCAL, DEFINE_REG, DETACH, DEVICE, ENCDEC, END, ENDIF, EXPAND, FLIP, GEP, GROUP, IDIV, IF, INDEX, INS, LINEAR, LOAD, LUNIQUE, MOD, MSELECT, MSTACK, MULACC, MULTI, NOOP, OR, PAD, PARAM, PERMUTE, PROGRAM, PTRCAT, RANGE, REDUCE, REDUCE_AXIS, RESHAPE, REWRITE_ERROR, SHL, SHR, SHRINK, SINK, SOURCE, SPECIAL, STORE, THREEFRY, UNIQUE, UNROLL, VCONST, VECTORIZE, WHERE, WMMA, and XOR.

When population_size is large enough, generation 0 is seeded with at least one instance of each supported primitive op.

About

Genetic programming on tinygrad

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages