botorch-rust v0.2.0

Second release of botorch-rust: a Rust implementation of BoTorch with the same Python API (import botorch_rs instead of import botorch), tested for equivalence against BoTorch.

Install

pip install --upgrade botorch-rust

Prebuilt abi3 wheels for CPython 3.11 and later, tested on 3.11, 3.14 and 3.15 (Windows x86_64, macOS arm64 and x86_64, Linux x86_64 and aarch64; no Rust toolchain needed), on PyPI and attached below. NumPy and PyTorch remain optional.

What is new

Coverage grew from 222 to 398 of BoTorch's 776 public classes and functions (51%) (coverage report, roadmap):

• Single-objective acquisition (details):
  • the knowledge gradient (also multi-fidelity), qNEI, qLCB and qPSD;
  • analytic noisy EI and LogNEI, LogProbabilityOfFeasibility, ScalarizedPosteriorMean and qAnalyticPI;
  • max-value entropy search and GIBBON (also multi-fidelity), joint and predictive entropy search;
  • Thompson sampling and qNegIntegratedPosteriorVariance;
  • the penalized, proximal, prior-guided and fixed-feature wrappers.
• Kernels, priors and constraints (details):
  • kernel trees (sums, products, scales, active_dims) in SingleTaskGP;
  • GPyTorch's RBF, Matérn, RQ, periodic, linear, polynomial, cosine and index kernels;
  • BoTorch's own kernels;
  • Normal, uniform, half-Cauchy, half-normal, LKJ, Beta and smoothed-box priors.
• Input and outcome transforms (details):
  • all 35 names of botorch.models.transforms, including the learned Warp and LearnedFeatureImputation (fitted jointly with the GP);
  • Log, Power, Bilog and StratifiedStandardize, with TransformedPosterior.
• Models:
  • multi-task: MultiTaskGP and KroneckerMultiTaskGP;
  • the rest: SingleTaskMultiFidelityGP, MixedSingleTaskGP, the MAP SAAS models, EnsembleModel, AffineDeterministicModel, FixedSingleSampleModel, ModelList / ModelDict and OrthogonalAdditiveGP;
  • subset_output and construct_inputs.
• Optimization and candidate generation (details):
  • optimize_acqf_mixed (also alternating), optimize_acqf_discrete (also with local search), optimize_acqf_cyclic and optimize_acqf_homotopy;
  • gen_candidates_torch, sample_around_best, post_processing_func and timeout_sec;
  • MaxPosteriorSampling and BoltzmannSampling;
  • fitting with torch optimizers, and batch_cross_validation.

Behaviour changes

• Fits at extreme hyperparameters now take BoTorch's path (details):
  • botorch-rust reproduces GPyTorch's NaN and infinite losses, so L-BFGS-B backtracks as in BoTorch. Before, some fits raised InvalidArgumentError or ended in a different optimum.
  • Cholesky factorizations try jitter up to 1e-3, as BoTorch configures linear_operator.
• Gradients exactly at clamp bounds follow PyTorch ≥ 2.14 (pytorch#191142) (details):
  • a scalar bound gives 0 at equality, and ties split 0.5/0.5;
  • with older PyTorch, BoTorch's gradients differ from botorch-rust's only at exact bounds.
• MultiTaskGP and LCEMGP raise UnsupportedError for outcome transforms other than a plain Standardize. Before, they silently treated them as one.
• Normalize.indices is a list.

Equivalence with BoTorch

The suite has 396 Rust tests and about 3,500 Python tests. It runs against BoTorch main and the latest release, under PyTorch 2.13 and 2.14.1.

Fits are compared in two stages:

1. deterministic checks on 16 twins (training inputs perturbed at rounding level);
2. if those fail, a statistical comparison with BoTorch's outcomes on 64 twins. It fails by chance with probability at most 0.2%.

The comparisons turned up bugs in BoTorch, GPyTorch and linear_operator, reported upstream:

• BoTorch: meta-pytorch/botorch#3427 and #3430–#3450;
• GPyTorch: cornellius-gp/gpytorch#2804–#2809;
• linear_operator: cornellius-gp/linear_operator#149.

Each docs page above lists the bugs it found.

Speed

Unchanged from 0.1.0; see the benchmarks.