jaxbo: Bayesian optimization in JAX

A maintained, modern-JAX Bayesian optimization library: a small stable core (exact GP regression, 5 kernels, the classic acquisition functions plus a batched score_candidates helper, input priors, L-BFGS training) with optional research extras (MCMC inference, multifidelity and manifold GPs, weighted sampling) that install on demand and are never paid for otherwise.

This is a maintained fork of PredictiveIntelligenceLab/JAX-BO; see the acknowledgment below.

Install

pip install jaxbo

# The quickstart below needs 0.2.0 (score_candidates, extras). Until it is
# on PyPI, install from the repo:
# pip install "jaxbo @ git+https://github.com/ricardogr07/JAX-BO"

The core installs with exactly 4 dependencies: jax, jaxlib, numpy, scipy. Everything else lives behind optional extras:

import jaxbo never imports an extra's dependencies; importing an extra without them raises an ImportError naming the pip install jaxbo[extra] fix.

Supported versions

The package pins jax>=0.6,<0.11 (matching jaxlib) and requires Python 3.10 or newer. CI tests every lane below at the jax floor and the newest jax the pin allows:

Quickstart: 60 seconds to the next point

Fit a GP to observations of an objective, then score a batch of candidates with expected improvement in one vmapped pass:

import jax.numpy as jnp
import numpy as np
from jax import random

from jaxbo.acquisitions import score_candidates
from jaxbo.gp import GP
from jaxbo.priors import uniform_prior
from jaxbo.utils import normalize


def f(x):
    return ((x - 1.5) ** 2).ravel()


# Domain and observations (raw domain)
lb, ub = jnp.array([-2.0]), jnp.array([3.0])
bounds = {"lb": lb, "ub": ub}
X = jnp.linspace(-2.0, 3.0, 8)[:, None]
y = f(X)

# GP expects an already normalized training batch
batch, norm_const = normalize(X, y, bounds)

gp = GP({"kernel": "RBF", "input_prior": uniform_prior(lb, ub), "criterion": "EI"})
params = gp.train(batch, random.PRNGKey(0), num_restarts=5)

# Score 256 raw-domain candidates in one vmapped pass and pick the best
X_cand = np.linspace(-2.0, 3.0, 256)[:, None]
scores = score_candidates(
    gp, X_cand, params=params, batch=batch, bounds=bounds,
    best=float(np.min(batch["y"])),
)
x_next = X_cand[np.argmin(scores)]
print("next point to evaluate:", x_next)  # close to the true minimum at 1.5

Two contract points worth knowing before you swap in a real objective:

• Normalization: train consumes batch["X"] exactly as given, so pass the already normalized batch (utils.normalize). predict and score_candidates normalize raw-domain inputs internally against bounds. Normalized batch in, raw candidates in; mixing these up fails silently.
• Scores: lower is better for every acquisition in jaxbo.acquisitions (EI is returned negated), so the next point is X_cand[np.argmin(scores)].

Prefer a notebook? Launch the interactive tutorial on Colab:

Docker

A quickstart image with jaxbo[all], JupyterLab, and the examples/ notebooks is published to GHCR with each release, starting with 0.2.0:

docker run -p 8888:8888 ghcr.io/ricardogr07/jaxbo

Then open the printed http://127.0.0.1:8888/lab?token=... URL.

Project documentation

• CONTRIBUTING.md: dev setup (uv + committed lockfile), tox envs, the CI gate map, PR rules, and the benchmark delta rule
• SCOPE.md: the 0.2.0 revamp scope, architecture, and locked decisions
• benchmarks/: the performance harness and the "no optimization without a delta" rule
• docs/audits/: reposage audit baseline the revamp is measured against
• CHANGELOG.md: release notes, generated by release-please

Acknowledgment

This project is a fork of JAX-BO by the Predictive Intelligence Lab at the University of Pennsylvania (Paris Perdikaris, Yibo Yang, Mohamed Aziz Bhouri). The core model structure, kernels, and acquisition functions originate there; this fork modernizes the library (current jax/python support, core/extras packaging, tests, benchmarks, CI) and is maintained independently by Ricardo García Ramírez. It is not affiliated with the original authors.

If you use this library in your research, please cite the original work (see also CITATION.cff):

@software{jaxbo2020github,
  author = {Paris Perdikaris, Yibo Yang, Mohamed Aziz Bhouri},
  title = {{JAX-BO}: A Bayesian optimization library in {JAX}},
  url = {https://github.com/PredictiveIntelligenceLab/JAX-BO},
  version = {0.2},
  year = {2020},
}

License

Apache License 2.0. See LICENSE.