Boax: A Bayesian Optimization library for JAX.

Overview | Installation | Getting Started | Documentation

Boax is currently in early alpha and under active development!

Overview

Boax is a composable library of core components for Bayesian Optimization that is designed for flexibility. It comes with low-level interfaces for:

• Core capabilities (boax.core):
  • Common Distributions
  • Monte-Carlo Samplers
• Fitting a surrogate model to data (boax.prediction):
  • Model Functions
  • Objective Functions
• Constructing and optimizing acquisition functions (boax.optimization):
  • Acquisition Functions
  • Optimizer Functions

Installation

You can install the latest released version of Boax from PyPI via:

pip install boax

or you can install the latest development version from GitHub:

pip install git+https://github.com/Lando-L/boax.git

Basic Usage

Here is a basic example of using the Boax API for defining a Gaussian Process model, constructing an Acquisition function, and generating the next batch of data points to query. For more details check out the docs.

1. Defining a Gaussian Process model:

from boax.prediction import models

model = models.gaussian_process.exact(
  models.means.zero(),
  models.kernels.scaled(
    models.kernels.rbf(1.0), 0.5
  ),
  models.likelihoods.gaussian(1e-4),
  x_train,
  y_train,
)

2. Constructing an Acquisition function.

from jax import vmap
from boax.optimization import acquisitions

acqf = models.outcome_transformed(
  vmap(model),
  acquisitions.upper_confidence_bound(2.0)
)

3. Generating the next batch of data points to query.

from jax import numpy as jnp
from jax import random
from boax.core import distributions, samplers
from boax.optimization import optimizers

key = random.key(0)

batch_size, num_results, num_restarts = 1, 100, 10
bounds = jnp.array([[-1.0, 1.0]])

sampler = samplers.halton_uniform(
  distributions.uniform.uniform(bounds[:, 0], bounds[:, 1])
)

optimizer = optimizers.batch(
  optimizers.initializers.q_batch(
    acqf, sampler, batch_size, num_results, num_restarts,
  ),
  optimizers.solvers.scipy(
    acqf, bounds,  
  ),
)

next_x, value = optimizer(key)

Citing Boax

To cite Boax please use the citation:

@software{boax2023github,
  author = {Lando L{\"o}per},
  title = {{B}oax: A Bayesian Optimization library for {JAX}},
  url = {https://github.com/Lando-L/boax},
  version = {0.1.3},
  year = {2023},
}

In the above bibtex entry, the version number is intended to be that from boax/version.py, and the year corresponds to the project's open-source release.