Hyperoptax: parallel hyperparameter tuning with JAX

Hyperoptax is a small, functional hyperparameter-optimization library for JAX. It supports nested PyTree search spaces, batched candidate evaluation, Python and jax.lax.scan loops, and random, grid, and Gaussian-process Bayesian search.

Installation

uv pip install hyperoptax

Install the optional notebook dependencies with:

uv pip install "hyperoptax[notebooks]"

Hyperoptax requires Python 3.10 or newer. For accelerator-specific JAX wheels, follow the JAX installation guide.

Quick start

Every optimizer follows the same functional pattern: init returns an immutable state PyTree and an optimizer, and the objective has signature objective(key, params) -> scalar.

import jax
import jax.numpy as jnp

from hyperoptax import BayesianSearch, LinearSpace, LogSpace


def objective(key, params):
    noise = 0.01 * jax.random.normal(key)
    return (
        (jnp.log10(params["learning_rate"]) + 3.0) ** 2
        + (params["dropout"] - 0.2) ** 2
        + noise
    )


space = {
    "learning_rate": LogSpace(1e-5, 1e-1),
    "dropout": LinearSpace(0.0, 0.5),
}

state, optimizer = BayesianSearch.init(
    space,
    n_max=80,
    n_parallel=4,
    maximize=False,
)
state, (params_history, results_history) = optimizer.optimize_scan(
    state,
    jax.random.PRNGKey(0),
    objective,
    n_iterations=20,
)

print(optimizer.best_result(state))
print(optimizer.best_params(state))

Here n_iterations=20 and n_parallel=4 perform 80 objective evaluations. For buffered optimizers, n_max is the number of observations the state can store, not the number of iterations. Keep n_max divisible by n_parallel to avoid evaluating a partially storable final batch.

Choosing an optimizer

Optimizer	Best suited to	Important constraints
RandomSearch	Cheap baseline and broad exploration	Stateless; does not retain observations
GridSearch	Small, explicitly enumerated grids	Every leaf must be a numeric DiscreteSpace; make the grid size divisible by n_parallel
BayesianSearch	Expensive, smooth objectives	Fixed n_max buffer; GP cost grows with the buffer size

All spaces may be nested in arbitrary dict/list/tuple PyTrees. Available leaf types are LinearSpace, LogSpace, DiscreteSpace, QLinearSpace, and QLogSpace. DiscreteSpace values must be numeric because optimizers represent candidates as JAX arrays.

Random and grid search

from hyperoptax import DiscreteSpace, GridSearch, RandomSearch

# Random search: 10 batches of 8 independent samples.
random_state, random_search = RandomSearch.init(space, n_parallel=8)
random_state, random_history = random_search.optimize(
    random_state, jax.random.PRNGKey(1), objective, n_iterations=10
)

# Grid search: six points, evaluated two at a time.
grid_space = {
    "width": DiscreteSpace((64, 128, 256)),
    "depth": DiscreteSpace((2, 4)),
}
grid_state, grid_search = GridSearch.init(grid_space, n_parallel=2)

Pass shuffle=True and an explicit PRNG key to randomize grid traversal reproducibly:

grid_state, grid_search = GridSearch.init(
    grid_space,
    key=jax.random.PRNGKey(2),
    shuffle=True,
    n_parallel=2,
)

Python loop versus JAX scan

optimize() runs the outer loop in Python and returns lists indexed by iteration:

• params_history: list of PyTrees whose leaves have leading shape (n_parallel, ...).
• results_history: list of arrays with shape (n_parallel,).

optimize_scan() uses jax.lax.scan, requires a JAX-traceable objective, and returns stacked arrays:

• params_history: PyTree with leaf shape (n_iterations, n_parallel, ...).
• results_history: array with shape (n_iterations, n_parallel).

Both paths use jax.vmap to evaluate each batch. Use the lower-level get_next_params and update_state methods when evaluations must run in an external system.

JAX constraints

The usual JAX sharp bits apply:

1. The objective must accept a PRNG key and one unbatched parameter PyTree, and return a scalar.
2. Every batch is evaluated with jax.vmap; data-dependent Python control flow and side effects are not supported.
3. Strings and variable-shaped model structures cannot be represented as batched JAX hyperparameters. Encode categorical choices numerically and map them to host-side labels where needed.
4. Parameters that change evaluation shape or duration can interact poorly with parallel batches; use n_parallel=1 when evaluations cannot be synchronized.

Notebooks

See notebooks/ for grid/Bayesian search examples, design studies, high-dimensional behavior, performance comparisons, RL tuning, and GP visualization.

Contributing

git clone https://github.com/TheodoreWolf/hyperoptax
cd hyperoptax
uv pip install -e ".[all]"
uv run pytest -m "not timing"
uv run ruff check src tests
uv run ruff format --check src tests

The full suite includes optional timing tests; run uv run pytest when you want those as well. Please keep objectives reproducible by passing and splitting JAX PRNG keys explicitly.

Roadmap

• Add callbacks for logging and early stopping.
• Reuse GP kernel blocks instead of rebuilding the full matrix each iteration.
• Replace the fixed number of length-scale Adam steps with a convergence-aware stopping criterion.
• Add more optimizer families for mixed and cost-sensitive search spaces.

Citation

@misc{hyperoptax,
  author = {Theo Wolf},
  title = {{Hyperoptax}: Parallel hyperparameter tuning with JAX},
  year = {2025},
  url = {https://github.com/TheodoreWolf/hyperoptax}
}