dl4bi 0.1.3

Deep Learning for Bayesian Inference (DL4BI)

Deep Learning for Bayesian Inference (dl4bi)

Install

Install with the appropriate command. If JAX isn't installed already, we recommend using one of the dl4bi[<jax-version>] installs.

pip install dl4bi # dl4bi
pip install dl4bi[cpu] # dl4bi + jax for CPU
pip install dl4bi[cuda12] # dl4bi + jax for CUDA-12
pip install dl4bi[cuda13] # dl4bi + jax for CUDA-13
pip install dl4bi[benchmarks,cpu] # benchmark deps + jax for CPU
pip install dl4bi[benchmarks,cuda12] # benchmark deps + jax for CUDA-12
pip install dl4bi[benchmarks,cuda13] # benchmark deps + jax for CUDA-13

The benchmarks extra installs the additional packages used by the benchmark scripts, especially under benchmarks/meta_learning and benchmarks/vae.

View Documentation (Locally)

git clone [e-mail]:MLGlobalHealth/dl4bi.git
cd dl4bi
uv run --with pdoc pdoc --docformat google --math dl4bi

Cite

If you're using this package or some of its code, please cite the relevant paper(s):

@inproceedings{jenson2026scalable,
  title = {Scalable Spatiotemporal Inference with Biased Scan Attention Transformer Neural Processes},
  author = {Jenson, Daniel and Navott, Jhonathan and Grynfelder, Piotr and Zhang, Mengyan and Sharma, Makkunda and Semenova, Elizaveta and Flaxman, Seth},
  booktitle = {Proceedings of the 29th International Conference on Artificial Intelligence and Statistics},
  series = {Proceedings of Machine Learning Research},
  volume = {300},
  address = {Tangier, Morocco},
  year = {2026},
  publisher = {PMLR}
}

@inproceedings{navott2026deeprv,
  title = {{DeepRV}: Accelerating Spatiotemporal Inference with Pre-trained Neural Priors},
  author = {Navott, Jhonathan and Jenson, Daniel and Flaxman, Seth and Semenova, Elizaveta},
  booktitle = {Proceedings of the 29th International Conference on Artificial Intelligence and Statistics},
  series = {Proceedings of Machine Learning Research},
  volume = {300},
  address = {Tangier, Morocco},
  year = {2026},
  publisher = {PMLR}
}

Benchmarks are available for BSA-TNP here and for DeepRV here.

Development Setup

• Install uv: curl -LsSf https://astral.sh/uv/install.sh | sh
• Clone the repository and cd into it: git clone [e-mail]:MLGlobalHealth/dl4bi.git && cd dl4bi
• Install the latest Python with uv: uv python install
• Sync the project environment:
  • CPU JAX: uv sync --extra cpu
  • CUDA 12 JAX: uv sync --extra cuda12
  • CUDA 13 JAX: uv sync --extra cuda13
  • Benchmark deps + CPU JAX: uv sync --extra benchmarks --extra cpu
  • Benchmark deps + CUDA 12 JAX: uv sync --extra benchmarks --extra cuda12
  • Benchmark deps + CUDA 13 JAX: uv sync --extra benchmarks --extra cuda13
• uv sync creates .venv, installs the project in editable mode, includes the default dev dependency group, and picks a Python interpreter compatible with the project's requires-python
• Before making changes, install the shared development hooks: uv run pre-commit install --install-hooks
• Verify the hook setup once per clone with: uv run pre-commit run --all-files
• Keep the hooks installed for local development; commits on main run pytest -q tests through the shared pre-commit setup
• Run project commands through uv, e.g. uv run pytest or uv run python gp.py
• If you want to activate the virtualenv directly, use source .venv/bin/activate

Build and Publish to PyPI

Create a local .env file with the publish tokens:

TEST_PYPI_TOKEN=pypi-...
PYPI_TOKEN=pypi-...

Run the release helper from a clean main checkout:

uv run python scripts/release.py .env "AISTATS 2026"

The helper bumps the patch version, commits and tags v<version> <message>, rebuilds dist/, publishes to TestPyPI and PyPI, pushes main and the tag, and smoke-tests the published install targets.