Getting started

Sbijax is a Python package for neural simulation-based inference and approximate Bayesian computation. Here we demonstrate its core functionality using a simple Gaussian model.

import jax
import sbijax
%matplotlib inline
import matplotlib.pyplot as plt

import matplotlib.pyplot as plt

def plot_posterior_hist(samples, key="theta", bins=40):
    theta = samples[key].reshape(-1, samples[key].shape[-1])
    d = theta.shape[-1]
    fig, axes = plt.subplots(1, d, figsize=(3 * d, 3))
    axes = [axes] if d == 1 else list(axes)
    for i, ax in enumerate(axes):
        ax.hist(theta[:, i], bins=bins, color="#700e01")
        ax.set_title(f"{key}[{i}]")
    fig.tight_layout(); return fig

Model definition

To do approximate inference using sbijax, a user first has to define a prior model and a simulator function which can be used to generate synthetic data. We will be using a simple bivariate Gaussian as an example with the following generative model:

\begin{align} \mu &\sim \mathcal{N}_2(0, I)\\ \sigma &\sim \mathcal{N}^+(1)\\ y & \sim \mathcal{N}_2(\mu, \sigma^2 I) \end{align}

Using TensorFlow Probability, the prior model and simulator are implemented like this:

from jax import numpy as jnp, random as jr
from tensorflow_probability.substrates.jax import distributions as tfd

def prior_fn():
    prior = tfd.JointDistributionNamed(dict(
        mean=tfd.Normal(jnp.zeros(2), 1.0),
        scale=tfd.HalfNormal(jnp.ones(1)),
    ), batch_ndims=0)
    return prior


def simulator_fn(seed: jr.PRNGKey, theta: dict[str, jax.Array]):
    p = tfd.Normal(jnp.zeros_like(theta["mean"]), 1.0)
    y = theta["mean"] + theta["scale"] * p.sample(seed=seed)
    return y

Algorithm definition

Having defined a model of interest, i.e., the prior and simulator functions, we construct an inferential method. Here, we use flow matching posterior estimation (FMPE) which is an algorithm that directly aims to learn an approximation to the posterior distribution (without intermediate steps like neural likelihood or neural likelihood-ratio estimation methods).

FMPE requires a continuous normalizing flow as a neural network model. We can construct it using the functionality provided by sbijax.

from sbijax.nn import make_cnf

n_dim_theta = 3
n_layers, hidden_size = 2, 128
neural_network = make_cnf(n_dim_theta, n_layers, hidden_size)

The neural network model in FMPE targets the posterior distribution and hence needs to target a space with its dimensionaliy. One can either find the dimensionality by inspecting the prior model (above, two for the mean and one for the scale), or use one of JAX internal methods, ravel_pytree, for this:

prior_draw = prior_fn().sample(seed=jr.PRNGKey(1))
prior_draw

{'scale': Array([0.9881485], dtype=float32),
 'mean': Array([ 1.2772286 , -0.66140693], dtype=float32)}

len(jax.flatten_util.ravel_pytree(prior_draw)[0])

3

We can then construct the method itself. Every method is a factory function that returns a record of pure functions (an Estimator). It takes as arguments the prior distribution and the neural network.

from sbijax import fmpe, simulate

prior = prior_fn()
model = fmpe(prior, neural_network)

Training and inference

Inference is then as easy as simulating some data, fitting the data to the model, a sampling from the approximate posterior. The data set is a dictionary of dictionaries (a PyTree in JAX lingo). It contains samples for the simulator function, called y, and parameter samples from the prior model, called theta.

data = simulate(
    jr.PRNGKey(1),
    prior,
    simulator_fn,
    n=10_000,
)
data

{'y': Array([[ 0.07692909,  0.7882271 ],
        [-1.2418504 , -0.25333643],
        [ 1.1943562 , -2.124853  ],
        ...,
        [ 1.3316323 ,  0.5488601 ],
        [ 4.862982  , -4.1227694 ],
        [-0.00955033,  0.989019  ]], dtype=float32),
 'theta': {'scale': Array([[0.9232396 ],
         [0.36471593],
         [0.6795394 ],
         ...,
         [0.11454558],
         [1.260745  ],
         [0.5012804 ]], dtype=float32),
  'mean': Array([[ 0.30212572,  0.67478853],
         [-0.963459  ,  0.086253  ],
         [ 0.39044896, -2.2378268 ],
         ...,
         [ 1.3652769 ,  0.6302657 ],
         [ 2.7543025 , -1.9100804 ],
         [-0.16545922,  0.6475435 ]], dtype=float32)}}

We then fit the model using the typical flow matching loss.

params, info = model.fit(jr.PRNGKey(2), data, n_iter=100, n_early_stopping_patience=20)

 92%|██████████████████████████▋  | 92/100 [00:56<00:04,  1.63it/s]

Finally, we sample from the posterior distribution for a specific observation \(y_{\text{obs}}\).

y_obs = jnp.array([-1.0, 1.0])
samples, _ = model.sample(
    jr.PRNGKey(3), params, y_obs, n_samples=1_000
)

Visualization

Sbijax provides basic functionality to analyse posterior draws. We show some visualizations below.

plt.plot(info.losses)
plt.xlabel("step")
plt.ylabel("loss")
plt.show()

plot_posterior_hist(samples, "mean")
plt.show()

Session info

import session_info

session_info.show(html=False)

-----
haiku                       0.0.16
jax                         0.10.2
jaxlib                      0.10.2
matplotlib                  3.11.0
sbijax                      0.4.0
session_info                v1.0.1
tensorflow_probability      0.26.0-dev20260318
-----
IPython             9.11.0
jupyter_client      8.8.0
jupyter_core        5.9.1
jupyterlab          4.5.6
notebook            7.5.5
-----
Python 3.12.10 (main, May 30 2025, 05:53:56) [Clang 20.1.4 ]
macOS-26.2-arm64-arm-64bit
-----
Session information updated at 2026-07-03 19:06