JAX-BO: Bayesian Optimization Library Maintenance

mlscientific-pythondeveloper-tooling

Modernized an academic JAX Bayesian-optimization library into an installable PyPI package

Now pip install jaxbo resolves to a real, current release

Situation

A JAX Bayesian-optimization library from the Predictive Intelligence Lab had stopped working on modern Python and JAX: the manifest claimed Python 3.6 compatibility but the package no longer installed, putting the research effectively out of reach.

Task

Fork it, restore modern Python/JAX compatibility, and ship it as a properly packaged, tested, automatically released PyPI library; credit the upstream research, claim only the distribution engineering.

Action

Forked the upstream library and restored Python and JAX compatibility; the manifest claimed Python 3.6, which was never real, and the actual working floor is ~3.10.
Packaged it for PyPI with Black and Ruff linting, plus CI on Python 3.10 and 3.12.
Wired up automated releases: release-please reads commits and bumps the version, a tag triggers the publish job, and a PyPI OIDC Trusted Publisher authenticates by identity instead of a stored token. Left the Gaussian-process and acquisition math credited to the upstream authors throughout.

How it works

Result

pip install jaxbo works today: two releases on PyPI (0.1.1 and 0.1.2, July 2025), shipped automatically via release-please versioning and a PyPI OIDC Trusted Publisher, with CI tests on Python 3.10 & 3.12 and Black + Ruff linting. An academic clone-and-run repo is now an installable, modern-Python package.

A chart showing a few tested points, a predicted curve, a shaded uncertainty band that is wide where there is no data, and an orange marker pointing to the next setting to try: high and still uncertain. — What Bayesian optimization does in plain terms: when each test is expensive, model both the predicted outcome and how unsure that prediction is, then test where the payoff is most uncertain-but-promising, homing in on the best setting in far fewer tries than a brute-force sweep.

Learning

In a JAX Gaussian process the hard part isn't differentiating the marginal likelihood: it's keeping that gradient finite. The closed form and its derivative have different domains of numerical safety, so the engineering goes into defending the second one: using a vector-Jacobian product instead of value-and-grad to dodge NaNs, epsilons under every square root, jitter on the Cholesky, and nanargmin across restarts because individual restarts will return NaN, and JIT will propagate that NaN silently through the whole loop before you see it. The surrogate is easy; the differentiable, JIT-able, multi-restart surrogate is where the work lives.

Tech Stack

PythonJAXGaussian ProcessesPyPIGitHub ActionsBlack / Ruff

Services

ML Research ToolingBayesian OptimizationScientific PythonPackage Maintenance

Status

On PyPI