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.
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_fitnessThe current recommended way to install tinygp is from source.
# Python3 installation
python3 -m pip install git+https://github.com/yippiez/tinygp.git
# For uv
uv pip install git+https://github.com/yippiez/tinygp.gitgit 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 .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/automatic_kernel_synthesis.pyshows GP-based kernel optimization starting from a tinygrad-defined reference kernel.
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.