👋 Welcome to sbijax!

👋 Welcome to `sbijax`!#

Simulation-based inference in JAX

Sbijax is a Python library for neural simulation-based inference and approximate Bayesian computation using JAX. It implements recent methods, such as Simulated Annealing ABC, Surjective Neural Likelihood Estimation, Neural Approximate Sufficient Statistics or Neural Posterior Score Estimation, as well as calibration and convergence diagnostics.

Caution

⚠️ As per the LICENSE file, there is no warranty whatsoever for this free software tool. If you discover bugs, please report them.

Quickstart#

Sbijax implements a low-level, functional API in the idiom of dm-haiku and blackjax: every method is a factory returning a tuple of pure functions. All a user needs to define is a prior function, a simulator function and an inferential algorithm. For example, you can define a neural likelihood estimation method and generate posterior samples like this:

from jax import numpy as jnp, random as jr
from sbijax import nle, train, sample, simulate
from sbijax.mcmc import make_sampler, nuts
from sbijax.nn import make_maf
from tensorflow_probability.substrates.jax import distributions as tfd

prior = tfd.JointDistributionNamed(dict(
    theta=tfd.Normal(jnp.zeros(2), jnp.ones(2))
), batch_ndims=0)

def simulator_fn(seed, theta):
    p = tfd.Normal(jnp.zeros_like(theta["theta"]), 0.1)
    y = theta["theta"] + p.sample(seed=seed)
    return y

estimator = nle(make_maf(2))

y_observed = jnp.array([-1.0, 1.0])
data = simulate(jr.key(1), prior, simulator_fn, n=10_000)
params, info = train(jr.key(2), estimator, data)
samples, _ = sample(
    jr.key(3), estimator, params, y_observed,
    sampler=make_sampler(nuts, prior=prior),
)

Workflow#

Every method in sbijax takes the same two user-supplied inputs:

a prior: needs a .sample(seed=rng_key) method returning a pytree, a batched form .sample(seed=rng_key, sample_shape=(n,)), and a .log_prob(theta) method accepting that same pytree structure. A pytree here just means a (possibly nested) dict of arrays, e.g. what the prior in the example above returns from .sample(...):
```
{"theta": Array([0.62, 0.84])}
```
Nothing checks that the prior is literally a tensorflow_probability.substrates.jax (tfd) distribution, but every example uses a tfd.JointDistributionNamed, which gives you both methods for free
a simulator: a plain function (seed, theta) -> y. Use those exact argument names (or **kwargs) — ABC methods (sabc/smcabc) call it by keyword internally. simulate/run_sequential always call it with a batched theta (a leading axis of size n), so it needs to handle a batch of parameter draws, not a single one

Caution

Each distribution in the prior must be vector-valued, even one that describes a single parameter: write tfd.Normal(jnp.zeros(1), 1.0), not tfd.Normal(0.0, 1.0). One draw then has shape (1,) and n draws have shape (n, 1).

The simulator must return y as a matrix of shape (n, d), even for one-dimensional data, i.e., (n, 1) instead of (n,).

From these two, the same pipeline applies to every neural estimator (NLE, NPE, FMPE, NPSE, NRE, SNLE):

        %%{init: {"flowchart": {"curve": "basis", "nodeSpacing": 60, "rankSpacing": 60, "htmlLabels": true, }, "themeVariables": {"fontFamily": "Helvetica, Arial, sans-serif", "fontSize": "20px", "nodePadding": "20px"}}}%%
flowchart LR
    P([prior]) -->|simulate| D[(data)]
    M([simulator]) -->|simulate| D
    D -->|"nle/npe/fmpe/..."| O([objective])
    O -->|train| T([params])
    T -->|sample| R([posterior])

    classDef input fill:#e8eef7,stroke:#5b7fa6,stroke-width:1.5px,color:#1c2b3a,font-weight:600;
    classDef data fill:#fbf3e3,stroke:#c99a3c,stroke-width:1.5px,color:#3a2f1c,font-weight:600;
    classDef obj fill:#eaf3ea,stroke:#4c8c5a,stroke-width:1.5px,color:#1c3a22,font-weight:600;
    classDef out fill:#f6e8ee,stroke:#a65b82,stroke-width:1.5px,color:#3a1c2b,font-weight:600;
    class P,M input
    class D data
    class O obj
    class T,R out

simulate(rng_key, prior, simulator, n) draws n prior/simulation pairs.
A factory (nle, npe, fmpe, …) wraps a neural network into an ObjectiveFns record of pure functions.
train(rng_key, objective, data) fits it, returning params.
sample(rng_key, objective, params, observable) draws posterior samples (MCMC-based methods also take sampler=make_sampler(nuts, prior=prior)).

Implemented methods#

The table lists each method with its factory function and the helper that builds a suitable network. ABC methods take the prior and simulator directly and need no network.

Method	Factory	Network helper	Reference
Sequential Monte Carlo ABC (SMC-ABC)	`smcabc()`	none, takes `summary_fn` and `distance_fn`	Beaumont et al. [2009]
Simulated annealing ABC (SABC)	`sabc()`	none, optional `summary_fn` and `distance_fn`	Albert et al. [2025]
Neural likelihood estimation (NLE)	`nle()`	`make_maf()`, `make_spf()`, `make_mdn()`	Papamakarios et al. [2019]
Surjective neural likelihood estimation (SNLE)	`snle()`	`make_maf()` or `make_spf()` with `n_layer_dimensions`	Dirmeier et al. [2023]
Neural posterior estimation (NPE)	`npe()`	`make_maf()`, `make_spf()`, `make_mdn()`	Greenberg et al. [2019]
Contrastive neural ratio estimation (NRE)	`nre()`	`make_mlp()`, `make_resnet()`	Miller et al. [2022]
Flow matching posterior estimation (FMPE)	`fmpe()`	`make_cnf()`	Wildberger et al. [2023]
Neural posterior score estimation (NPSE)	`npse()`	`sbijax.experimental.nn.make_score_model()`	Sharrock et al. [2024]
All-in-one posterior estimation (AIO)	`sbijax.experimental.aio()`	`sbijax.experimental.nn.make_simformer_based_score_model()`	Gloeckler et al. [2024]
Consistency model posterior estimation (CMPE)	`sbijax.experimental.cmpe()`	`make_cm()`	Schmitt et al. [2023]
Neural approximate sufficient statistics (NASS)	`nass()`	`make_nass_net()`	Chen et al. [2021]
Neural approximate slice sufficient statistics (NASSS)	`nasss()`	`make_nasss_net()`	Chen et al. [2023]

NASS and NASSS learn summary statistics, not a posterior. Pass the fitted summary network to summarized_estimator() to train a neural estimator on the summaries, or use its summarize_fn as the summary_fn of an ABC method. The network helpers are defaults: a factory accepts any network with the methods listed in its docstring (see Using Flax linen networks). Methods in sbijax.experimental may change or be removed.

Installation#

You can install sbijax from PyPI using:

pip install sbijax

To install the latest GitHub <RELEASE>, just call the following on the command line:

pip install git+https://github.com/dirmeier/sbijax@<RELEASE>

See also the installation instructions for JAX, if you plan to use sbijax on GPU/TPU.

Contributing and Support#

If you have questions, encounter problems, or need support with this software, please use the following channels:

Questions & Discussions: For general questions, usage help, or architectural discussions, please open a new thread in our GitHub Discussions tab.
Bug Reports & Feature Requests: To report a bug, software problem, or suggest a new feature, please submit an issue via our GitHub Issue Tracker. Please check existing issues before opening a new one to ensure it hasn’t already been reported.

Code contributions in the form of pull requests are more than welcome. A good way to start is to check out issues labelled good first issue. If you are unsure, if starting to work on a PR makes sense, feel free to open an issue or discussion thread.

In order to contribute:

Clone sbijax and install uv from here.
Install all dependencies using uv sync --all-groups.

Install the Git hooks:

uv run pre-commit install -t pre-commit -t commit-msg

Create a new branch locally git checkout -b feature/my-new-feature or git checkout -b issue/fixes-bug.
Implement your contribution and ideally a test case.
Check your work (see below).
Submit a PR 🙂

Development commands#

The project uses uv for everything:

uv sync --all-groups
uv run pytest # runs only fast tests
uv run pytest -m slow # runs all tests
uv run ruff check sbijax examples
uv run ruff check --fix sbijax examples
uv run ruff format sbijax examples
uv run mypy sbijax examples
uv run pre-commit run --all-files

License#

sbijax is licensed under the Apache 2.0 License.