Quick start#

bde implements Bayesian Deep Ensembles that plug directly into scikit-learn pipelines while running training and sampling in JAX. This page walks through installing the package, setting up the environment and information regarding JAX device count.

Installation#

To install the latest release from PyPI, run:

pip install sklearn-contrib-bde

To install the latest development version from GitHub, run:

pip install git+https://github.com/scikit-learn-contrib/bde.git

If you prefer reproducible environments with GPU or multi-device CPU support, consider using pixi as described in User Guide. The repo already ships with a pixi.toml that pins the required versions of JAX, scikit-learn, and the CUDA toolchain when available.

Set the JAX device count#

When you run the estimators outside of the examples shipped in examples/, ensure JAX can see enough host devices. The environment variable below makes CPU-only runs allocate eight virtual devices; tweak the value to match your hardware:

export XLA_FLAGS="--xla_force_host_platform_device_count=<desired_devices_to_be_used>"

You can set it once per shell session or inject it programmatically before importing JAX in your Python scripts.

Work with the development environment#

Once you clone the repository, bootstrap the dev dependencies and tooling with:

pixi install

You can then run the standard tasks:

pixi run lint for style checks (ruff and black)
pixi run test for the pytest suite, including scikit-learn compatibility checks in tests/test_common.py
pixi run build-doc to render this documentation locally with Sphinx and sphinx-gallery examples

If you prefer calling Python tooling directly, enter the environment shell:

pixi shell -e dev

From there pytest, ruff and sphinx-build are available on PATH.

Next steps#

Dive into User Guide for a walkthrough of the package internals and configuration knobs.
Explore the auto-generated API reference in API Reference to see every public class, function, and dataclass.
Run the gallery in Examples to compare regression and classification behaviours on real-world datasets!