Using Hyperoptax

Using Hyperoptax#

Installation#

uv pip install hyperoptax

Hyperoptax requires Python 3.10 or newer. Install accelerator-specific JAX wheels by following the JAX installation guide.

Choosing an optimizer#

Optimizer

Best suited to

Constraints

RandomSearch

Cheap baseline and broad exploration

Stateless; does not retain observations.

GridSearch

Small, explicitly enumerated grids

Every leaf must be a numeric DiscreteSpace. The grid size should be divisible by n_parallel.

BayesianSearch

Expensive, smooth objectives

Stores observations in a fixed n_max buffer; GP work grows with the buffer size.

Iterations, batches, and buffers#

n_parallel is the number of configurations evaluated in each optimizer iteration. Therefore n_iterations * n_parallel is the number of objective evaluations.

Bayesian search uses fixed-size buffers so its state has static JAX shapes. n_max is the number of observations that fit in the buffer, not the number of iterations. When n_iterations is omitted, Bayesian search runs until the buffer is full. Keep n_max divisible by n_parallel so the last batch fits exactly.

Grid search builds the full Cartesian product during init. If the grid size is not divisible by n_parallel, its incomplete final batch is omitted.

Loop APIs#

optimize() runs the outer loop in Python. It returns params_history and results_history as lists indexed by iteration. The parameter PyTrees have a leading n_parallel axis and each result array has shape (n_parallel,).

optimize_scan() uses jax.lax.scan and requires a JAX-traceable objective. It returns stacked parameter leaves with shape (n_iterations, n_parallel, ...) and results with shape (n_iterations, n_parallel).

Both APIs evaluate the batch with jax.vmap. The objective must accept one unbatched configuration and return a scalar.

External evaluation loop#

Use get_next_params() and update_state() directly when a job scheduler, service, or other external system evaluates configurations.

import jax

key = jax.random.PRNGKey(0)
previous_params = None
previous_results = None

for _ in range(20):
    key, ask_key, eval_key, tell_key = jax.random.split(key, 4)
    params = optimizer.get_next_params(
        state, ask_key, previous_params, previous_results
    )
    eval_keys = jax.random.split(eval_key, optimizer.n_parallel)
    results = jax.vmap(objective)(eval_keys, params)
    state = optimizer.update_state(state, tell_key, results, params)
    previous_params, previous_results = params, results

JAX constraints#

  • Search-space leaves and objective results must be representable as JAX arrays. String categories are not supported; encode them numerically.

  • Data-dependent Python control flow, variable-shaped outputs, and side effects are not compatible with optimize_scan.

  • Parameters that change evaluation shape or duration may not batch cleanly. Use n_parallel=1 when evaluations cannot be synchronized.

Contents