sbijax

https://sbijax.readthedocs.io/en/latest/
Admin
sbijax is a Python library for neural simulation-based inference and approximate Bayesian...
Tokens:91,472
Snippets:445
Trust Score:9.6
Update:1 month ago
Show doc for...
Tokens:
Raw
### Install sbijax from PyPI

Source: https://sbijax.readthedocs.io/en/latest/index.html

Installs the latest released version of sbijax from PyPI using pip.

```bash
pip install sbijax
```

--------------------------------

### Install sbijax from GitHub release

Source: https://sbijax.readthedocs.io/en/latest/index.html

Installs a specific GitHub release of sbijax by replacing <RELEASE> with the desired release tag.

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

--------------------------------

### Import libraries for sbijax examples

Source: https://sbijax.readthedocs.io/en/latest/notebooks/examples.html

Imports the libraries used throughout the examples: jax, optax, os, sbijax, seaborn, matplotlib, and TensorFlow Probability's JAX substrate. Enables inline plotting with %matplotlib inline.

```python
import jax
import optax
import os
import sbijax
import seaborn as sns
%matplotlib inline
import matplotlib.pyplot as plt

from matplotlib.ticker import AutoLocator, MaxNLocator
from jax import numpy as jnp, random as jr
from tensorflow_probability.substrates.jax import distributions as tfd

```

--------------------------------

### sample_with_mala example

Source: https://sbijax.readthedocs.io/en/latest/api/sbijax.mcmc.html

Samples from a distribution using the MALA sampler. The example defines a prior and a log-probability function, then calls sample_with_mala with a random key and the partial log-probability.

```python
>>> import functools as ft
>>> from jax import numpy as jnp, random as jr
>>> from tensorflow_probability.substrates.jax import distributions as tfd
...
>>> prior = tfd.JointDistributionNamed(
...    dict(theta=tfd.Normal(jnp.zeros(2), 1.0))
... )
>>> def log_prob(theta, y):
...     lp_prior = prior.log_prob(theta)
...     lp_data = tfd.Normal(theta["theta"], 1.0).log_prob(y)
...     return jnp.sum(lp_data) + jnp.sum(lp_prior)
...
>>> prop_posterior_lp = ft.partial(log_prob, y=jnp.array([-1.0, 1.0]))
>>> samples = sample_with_mala(jr.key(0), prop_posterior_lp, prior)
```

--------------------------------

### Show session info with session_info.show(html=False)

Source: https://sbijax.readthedocs.io/en/latest/notebooks/eeg_data_example.html

Use this snippet to display the versions of installed packages and system information in a Jupyter notebook. The output shows the package versions, Python version, and platform details. The html=False argument makes the output plain text instead of HTML.

```python
import session_info

session_info.show(html=False)

```

```text
-----
arviz                       0.19.0
haiku                       0.0.12
jax                         0.4.31
jaxlib                      0.4.31
jrnmm                       0.1.0.post1
matplotlib                  3.9.2
mne                         1.8.0
moabb                       1.1.1
numpy                       1.26.4
optax                       0.2.4
pandas                      1.5.3
sbijax                      0.3.3
seaborn                     0.12.2
session_info                1.0.0
tensorflow_probability      0.25.0
-----
IPython             8.31.0
jupyter_client      8.6.3
jupyter_core        5.7.2
jupyterlab          4.3.1
-----
Python 3.11.7 (main, Dec  9 2023, 06:06:18) [Clang 14.0.3 (clang-1403.0.22.14.1)]
macOS-13.0.1-arm64-arm-64bit
-----
Session information updated at 2025-01-04 20:56

```

--------------------------------

### Fit summary network and wrap estimator with summarized_estimator

Source: https://sbijax.readthedocs.io/en/latest/_modules/sbijax/_src/inference/summary/_compose.html

Fit a summary network first, then wrap a downstream estimator with summarized_estimator so training and sampling operate on learned summaries. The example shows the required order: train the summary net, wrap the estimator, then train and sample from the wrapped estimator.

```python
sn = nass(make_nass_net(2, [64, 64]))
      sn_params, _ = train(key, sn, data)
      est = summarized_estimator(nle(make_maf(2)), sn, sn_params)
      params, info = train(key, est, data)           # trains on summaries
      samples, _ = est.sample_fn(key, params, y_observed, sampler=sampler)
```

--------------------------------

### Neural likelihood estimation example with sbijax

Source: https://sbijax.readthedocs.io/en/latest/index.html

Shows the low-level functional API of sbijax: define a prior and simulator, create an estimator with nle and a MAF network, simulate data, train, and sample posterior draws. Requires JAX, TensorFlow Probability (JAX backend), and sbijax.

```python
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),
)
```

--------------------------------

### Sampling with NUTS

Source: https://sbijax.readthedocs.io/en/latest/_modules/sbijax/_src/mcmc/nuts.html

Example usage of the sample_with_nuts function to sample from a distribution.

```python
>>> import functools as ft
>>> from jax import numpy as jnp, random as jr
>>> from tensorflow_probability.substrates.jax import distributions as tfd
...
>>> prior = tfd.JointDistributionNamed(
...    dict(theta=tfd.Normal(jnp.zeros(2), 1.0))
... )
>>> def log_prob(theta, y):
...     lp_prior = prior.log_prob(theta)
...     lp_data = tfd.Normal(theta["theta"], 1.0).log_prob(y)
...     return jnp.sum(lp_data) + jnp.sum(lp_prior)
...
>>> prop_posterior_lp = ft.partial(log_prob, y=jnp.array([-1.0, 1.0]))
>>> samples = sample_with_nuts(jr.key(0), prop_posterior_lp, prior)
```

--------------------------------

### Simulate data with sbijax.simulate and stack (0.4)

Source: https://sbijax.readthedocs.io/en/latest/migration.html

Shows the 0.3 vs 0.4 API for data simulation. In 0.4, use `simulate` with a prior and simulator function, and `stack` to append additional rounds from a proposal.

```python
# 0.3
data, _ = model.simulate_data(jr.PRNGKey(0), n_simulations=10_000)

# 0.4
from sbijax import simulate, stack
data = simulate(jr.key(0), prior, simulator_fn, n=10_000)
# append another round drawn from a proposal:
more = simulate(jr.key(1), prior, simulator_fn, proposal=proposal, n=10_000)
data = stack(data, more)

```

--------------------------------

### sample_with_imh example

Source: https://sbijax.readthedocs.io/en/latest/api/sbijax.mcmc.html

Draws samples using the independent Metropolis-Hastings sampler. The example defines a prior and a log-probability function, then calls sample_with_imh with a random key and the partial log-probability.

```python
>>> import functools as ft
>>> from jax import numpy as jnp, random as jr
>>> from tensorflow_probability.substrates.jax import distributions as tfd
...
>>> prior = tfd.JointDistributionNamed(
...    dict(theta=tfd.Normal(jnp.zeros(2), 1.0))
... )
>>> def log_prob(theta, y):
...     lp_prior = prior.log_prob(theta)
...     lp_data = tfd.Normal(theta["theta"], 1.0).log_prob(y)
...     return jnp.sum(lp_data) + jnp.sum(lp_prior)
...
>>> prop_posterior_lp = ft.partial(log_prob, y=jnp.array([-1.0, 1.0]))
>>> samples = sample_with_imh(jr.key(0), prop_posterior_lp, prior)
```

--------------------------------

### make_maf example in Python

Source: https://sbijax.readthedocs.io/en/latest/api/sbijax.nn.html

Example usage: creates an affine masked autoregressive flow with a custom dimension structure. The n_layer_dimensions argument determines which layers are dimensionality-preserving or -reducing.

```python
>>> neural_network = make_maf(10, n_layer_dimensions=(10, 10, 5, 5, 5))

```

--------------------------------

### Custom training loop with obj.train.init_fn and obj.train.step_fn

Source: https://sbijax.readthedocs.io/en/latest/custom_loops.html

Shows a full custom training loop using the low-level primitives. It flattens the prior pytree to a (n, d) array, builds a batch generator, initialises the TrainingState from a first batch, jits the step function, and runs 100 epochs with per-epoch key folding. The comment shows where to plug in your own validation or early stopping; the final params are exactly what sbijax.train() would have returned.

```python
import jax
import optax
from jax import random as jr
from jax._src.flatten_util import ravel_pytree
from sbijax import npe, simulate
from sbijax.nn import make_maf

obj = npe(make_maf(2))
optimizer = optax.adam(3e-4)

data = simulate(jr.key(0), prior, simulator, n=10_000)
# flatten the prior-pytree parameters to a (n, d) array
theta = jax.vmap(lambda t: ravel_pytree(t)[0])(data["theta"])
xy = {"y": data["y"], "theta": theta}

def get_batch(xy, batch_size, key):
    n = xy["y"].shape[0]
    perm = jr.permutation(key, n)
    for i in range(0, n - batch_size + 1, batch_size):
        idx = perm[i : i + batch_size]
        yield {k: v[idx] for k, v in xy.items()}

# build the carry from a first batch, then jit the step
first = next(get_batch(xy, 128, jr.key(1)))
state = obj.train.init_fn(optimizer, jr.key(2), first)
step = jax.jit(lambda rng, s, b: obj.train.step_fn(optimizer, rng, s, b))

key = jr.key(3)
for epoch in range(100):
    for batch in get_batch(xy, 128, jr.fold_in(key, epoch)):
        key, step_key = jr.split(key)
        metrics, state = step(step_key, state, batch)
    # plug in your own validation / early stopping, e.g.:
    # val = obj.train.eval_fn(jr.fold_in(key, epoch), state, val_batch)

params = state.params
```

--------------------------------

### Flatten prior draw to get dimensionality with ravel_pytree

Source: https://sbijax.readthedocs.io/en/latest/notebooks/getting_started.html

Uses JAX's flatten_util.ravel_pytree to flatten the prior sample and counts its length to get the total dimensionality of the parameter space.

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

--------------------------------

### sbijax._src.mcmc.irmh: sample_with_imh and IMH kernel setup

Source: https://sbijax.readthedocs.io/en/latest/_modules/sbijax/_src/mcmc/irmh.html

Defines the independent Metropolis-Hastings (IMH) sampler for sbijax. The public function `sample_with_imh` draws samples using BlackJAX's `bj.irmh` kernel, with defaults of 4 chains, 2,000 samples per chain, and 1,000 warmup samples. The module also provides internal helpers `_irmh_proposal_distribution` and `_mh_init`, and exposes the sampler as `imh = Kernel(init_fn=_mh_init)`. The docstring includes a usage example that builds a joint prior and a log-probability function, then calls `sample_with_imh` with a JAX random key.

```python
import blackjax as bj
import jax
from jax import random as jr

from sbijax._src.mcmc.sampler import Kernel
from sbijax._src.mcmc.util import run_blackjax


# ruff: noqa: PLR0913, D417


[docs]
def sample_with_imh(
  rng_key, lp, prior, *, n_chains=4, n_samples=2_000, n_warmup=1_000, **kwargs
):
  r"""Draw samples using the independent Metropolis-Hastings sampler.

  Args:
      rng_key: a jax random key
      lp: the logdensity you wish to sample from
      prior: a function that returns a prior sample
      n_chains: number of chains to sample
      n_samples: number of samples per chain
      n_warmup: number of samples to discard

  Examples:
      >>> import functools as ft
      >>> from jax import numpy as jnp, random as jr
      >>> from tensorflow_probability.substrates.jax import distributions as tfd
      ...
      >>> prior = tfd.JointDistributionNamed(
      ...    dict(theta=tfd.Normal(jnp.zeros(2), 1.0))
      ... )
      >>> def log_prob(theta, y):
      ...     lp_prior = prior.log_prob(theta)
      ...     lp_data = tfd.Normal(theta["theta"], 1.0).log_prob(y)
      ...     return jnp.sum(lp_data) + jnp.sum(lp_prior)
      ...
      >>> prop_posterior_lp = ft.partial(log_prob, y=jnp.array([-1.0, 1.0]))
      >>> samples = sample_with_imh(jr.key(0), prop_posterior_lp, prior)

  Returns:
      a tuple ``(samples, info)``: a named pytree with leaves of shape
      ``n_chains x (n_samples - n_warmup) x dim`` and an
      ``MCMCSampleInfo`` with the mean post-warmup acceptance rate
  """
  init_key, run_key = jr.split(rng_key)
  initial_positions = prior.sample(seed=init_key, sample_shape=(n_chains,))
  return run_blackjax(
    run_key,
    _mh_init,
    initial_positions,
    lp,
    n_chains=n_chains,
    n_samples=n_samples,
    n_warmup=n_warmup,
  )




# ruff: noqa: E731
def _irmh_proposal_distribution(initial_positions):
  shape = lambda p: () if p.ndim == 1 else (p.shape[-1],)

  def fn(rng_key):
    return {
      k: jax.random.normal(rng_key, shape=shape(v))
      for k, v in initial_positions.items()
    }

  return fn


# pylint: disable=missing-function-docstring,no-member
def _mh_init(_rng_key, initial_positions, lp):
  kernel = bj.irmh(lp, _irmh_proposal_distribution(initial_positions))
  initial_state = jax.vmap(kernel.init)(initial_positions)
  return initial_state, kernel.step


imh = Kernel(init_fn=_mh_init)
```

--------------------------------

### Initialize and train an estimator

Source: https://sbijax.readthedocs.io/en/latest/api/sbijax.html

Demonstrates the factory pattern for creating a neural likelihood estimator and the subsequent use of driver functions for training and sampling.

```python
est = nle(make_maf(2))
params, info = train(key, est, data, optimizer=optax.adam(3e-4))
samples, info = sample(
    key, est, params, y_observed, sampler=make_sampler(nuts, prior=prior)
)
```

--------------------------------

### Import libraries and set up matplotlib

Source: https://sbijax.readthedocs.io/en/latest/notebooks/getting_started.html

Imports JAX, Sbijax, and matplotlib, and enables inline plotting in Jupyter notebooks.

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

--------------------------------

### Sampling with RMH

Source: https://sbijax.readthedocs.io/en/latest/api/sbijax.mcmc.html

Example of using sample_with_rmh to sample from a distribution using the Rosenbluth-Metropolis-Hastings sampler.

```python
>>> import functools as ft
>>> from jax import numpy as jnp, random as jr
>>> from tensorflow_probability.substrates.jax import distributions as tfd
...
>>> prior = tfd.JointDistributionNamed(
...    dict(theta=tfd.Normal(jnp.zeros(2), 1.0))
... )
>>> def log_prob(theta, y):
...     lp_prior = prior.log_prob(theta)
...     lp_data = tfd.Normal(theta["theta"], 1.0).log_prob(y)
...     return jnp.sum(lp_data) + jnp.sum(lp_prior)
...
>>> prop_posterior_lp = ft.partial(log_prob, y=jnp.array([-1.0, 1.0]))
>>> samples = sample_with_rmh(jr.key(0), prop_posterior_lp, prior)
```

--------------------------------

### Simulate data with sbijax.simulate

Source: https://sbijax.readthedocs.io/en/latest/notebooks/more_detailed_intro.html

Simulates 10,000 data points from the prior and simulator using sbijax.simulate with a JAX random key. The returned data is a dictionary of dictionaries (a PyTree) containing 'y' samples and 'theta' parameter samples, as shown in the displayed output.

```python
from sbijax import simulate

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

```

```output
{'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)}}

```

--------------------------------

### Sample from likelihood/ratio posterior with sbijax.sample and make_sampler

Source: https://sbijax.readthedocs.io/en/latest/migration.html

For likelihood/ratio methods (`nle`/`nre`), the prior travels inside a sampler built with `sbijax.mcmc.make_sampler()`. This allows reusing one trained likelihood under different priors without retraining.

```python
from sbijax import sample
from sbijax.mcmc import make_sampler, nuts

samples, info = sample(
    jr.key(2), estimator, params, y_obs,
    sampler=make_sampler(nuts, prior=prior),
)

```

--------------------------------

### Sample from a distribution using sample_with_slice

Source: https://sbijax.readthedocs.io/en/latest/_modules/sbijax/_src/mcmc/slice.html

Example usage of the sample_with_slice function to perform MCMC sampling with a defined prior and log-probability function.

```python
>>> import functools as ft
>>> from jax import numpy as jnp, random as jr
>>> from tensorflow_probability.substrates.jax import distributions as tfd
...
>>> prior = tfd.JointDistributionNamed(
...     dict(
...         mean=tfd.Normal(jnp.zeros(2), 1.0),
...         std=tfd.HalfNormal(1.0)
...     ),
...     batch_ndims=0,
... )
>>> def log_prob(theta, y):
...     lp_prior = prior.log_prob(theta)
...     lp_data = tfd.Normal(theta["mean"], theta["std"]).log_prob(y)
...     return jnp.sum(lp_data) + jnp.sum(lp_prior)
...
>>> prop_posterior_lp = ft.partial(log_prob, y=jnp.array([-1.0, 1.0]))
>>> samples = sample_with_slice(jr.PRNGKey(0), prop_posterior_lp, prior)
```

--------------------------------

### ScoreModel sample method

Source: https://sbijax.readthedocs.io/en/latest/_modules/sbijax/_src/experimental/nn/make_score_network.html

Samples from the score model by integrating the reverse SDE from time_max to time_eps using RK45, starting from the base distribution.

```python
def sample(self, context, **kwargs):
    """Sample from the score model.

    Args:
        context: array of conditioning variables
    """
    theta_0 = self._base_distribution.sample(
      seed=hk.next_rng_key(), sample_shape=(context.shape[0],)
    )
    drift_and_diffusion_fn = get_forward_drift_and_diffusion_fn(
      self._sde, beta_min=self._beta_min, beta_max=self._beta_max
    )

    def sde_fn(time, theta_t):
      theta_t = theta_t.reshape(-1, self._n_dimension)
      time = jnp.repeat(time, theta_t.shape[0])
      drift, diffusion = drift_and_diffusion_fn(theta_t, time)
      score = self._score_net(
        inputs=theta_t, time=time, context=context, **kwargs
      )
      ret = drift - 0.5 * diffusion**2 * score
      return ret.reshape(-1)

    res = integrate.solve_ivp(
      sde_fn,
      (self._time_max, self._time_eps),
      theta_0.reshape(-1),
      rtol=1e-5,
      atol=1e-5,
      method="RK45",
    )

    ret = res.y[:, -1].reshape(-1, self._n_dimension)
    return ret
```

--------------------------------

### Sample posterior for closed eyes EEG

Source: https://sbijax.readthedocs.io/en/latest/notebooks/eeg_data_example.html

Samples the posterior for EEG data with closed eyes. Uses the same estimator and parameters as the synthetic example, but with the experimental closed-eyes observation.

```python
posterior_closed, _ = estim.sample_posterior(
    jr.PRNGKey(2),
    params,
    observable=summarize(eeg_data["y_experimental_closed"]),
    n_samples=10_000
)
```

--------------------------------

### Sample posterior with slice sampler (sbijax.mcmc.sample_with_slice)

Source: https://sbijax.readthedocs.io/en/latest/notebooks/examples.html

Sets up the slice sampler for MCMC sampling from the posterior. The likelihood and log-density functions define the model, and the sampling call uses 10 chains, 10,000 samples, and 5,000 warmup steps. The resulting samples are stored in `slice_samples`.

```python
from functools import partial
from jax import scipy as jsp
from sbijax.mcmc import sample_with_nuts, sample_with_slice
```

```python
def likelihood_fn(theta, y):
    mu = jnp.tile(theta[:2], 4)
    s1, s2 = theta[2] ** 2, theta[3] ** 2
    corr = s1 * s2 * jnp.tanh(theta[4])
    cov = jnp.array([[s1**2, corr], [corr, s2**2]])
    cov = jsp.linalg.block_diag(*[cov for _ in range(4)])
    p = tfd.MultivariateNormalFullCovariance(mu, cov)
    return p.log_prob(y)


def log_density_fn(theta, y):
    prior_lp = tfd.JointDistributionNamed(dict(
        theta=tfd.Uniform(jnp.full(5, -3.0), jnp.full(5, 3.0))
    )).log_prob(theta)
    likelihood_lp = likelihood_fn(theta, y)
    lp = jnp.sum(prior_lp) + jnp.sum(likelihood_lp)
    return lp
```

```python
log_density = partial(log_density_fn, y=y_obs)

def lp(theta):
    return log_density(theta["theta"])

slice_samples_raw, _ = sample_with_slice(
    jr.PRNGKey(0),
    lp,
    prior_fn(),
    n_chains=10,
    n_samples=10_000,
    n_warmup=5_000
)
slice_samples = slice_samples_raw["theta"]
```

--------------------------------

### Define prior_fn and simulator_fn for generative model

Source: https://sbijax.readthedocs.io/en/latest/notebooks/examples.html

Defines the generative model for the 'Simple Likelihood Complex Posterior' example. prior_fn returns a joint distribution with a 5-dimensional uniform prior on theta; simulator_fn generates 8-dimensional observations from theta using a correlated normal structure.

```python
def prior_fn():
    prior = tfd.JointDistributionNamed(dict(
        theta=tfd.Uniform(jnp.full(5, -3.0), jnp.full(5, 3.0))
    ), batch_ndims=0)
    return prior


def simulator_fn(seed, theta):
    theta = theta["theta"]
    theta = theta[:, None, :]
    us_key, noise_key = jr.split(seed)

    def _unpack_params(ps):
        m0 = ps[..., [0]]
        m1 = ps[..., [1]]
        s0 = ps[..., [2]] ** 2
        s1 = ps[..., [3]] ** 2
        r = jnp.tanh(ps[..., [4]])
        return m0, m1, s0, s1, r

    m0, m1, s0, s1, r = _unpack_params(theta)
    us = tfd.Normal(0.0, 1.0).sample(
        seed=us_key, sample_shape=(theta.shape[0], theta.shape[1], 4, 2)
    )
    xs = jnp.empty_like(us)
    xs = xs.at[:, :, :, 0].set(s0 * us[:, :, :, 0] + m0)
    y = xs.at[:, :, :, 1].set(
        s1 * (r * us[:, :, :, 0] + jnp.sqrt(1.0 - r**2) * us[:, :, :, 1]) + m1
    )
    y = y.reshape((*theta.shape[:1], 8))
    return y

```

--------------------------------

### Objective primitives: train.init_fn, train.step_fn, train.eval_fn, sample_fn

Source: https://sbijax.readthedocs.io/en/latest/custom_loops.html

Lists the bound pure functions exposed by a trainable objective. Use these low-level primitives when you need full control over training or sampling instead of the convenience drivers sbijax.train() and sbijax.sample(). TrainingState is an opaque carry; params is what you extract at the end.

```python
obj.train.init_fn(optimizer, rng, batch)        -> TrainingState(params, opt_state)
obj.train.step_fn(optimizer, rng, state, batch) -> (metrics, TrainingState)
obj.train.eval_fn(rng, state, batch)            -> metrics
obj.sample_fn(rng, params, observable, *, sampler=None) -> (samples, info)
```

--------------------------------

### sbijax.simulate(rng_key, prior, simulator, *, proposal=None, n)

Source: https://sbijax.readthedocs.io/en/latest/api/sbijax.html

Draws a dataset of parameters and observations. Parameters are drawn from the proposal if given, otherwise from the prior, and observations are produced by running the simulator on them.

```APIDOC
## sbijax.simulate(rng_key, prior, simulator, *, proposal=None, n)

### Description
Draw a dataset of parameters and observations. Parameters are drawn from `proposal` if given, otherwise from `prior`, and observations are produced by running `simulator` on them.

### Parameters
- **rng_key** – a jax random key
- **prior** – a distribution to draw parameters from when `proposal` is None
- **simulator** – a callable `(rng_key, theta) -> y`
- **proposal** – an optional callable `(rng_key, n) -> theta` used to draw parameters instead of the prior (e.g. a fitted posterior in a sequential round)
- **n** – the number of parameter/observation pairs to draw

### Returns
a dictionary with keys `y` and `theta`
```

--------------------------------

### Show session info with session_info

Source: https://sbijax.readthedocs.io/en/latest/notebooks/examples.html

Prints the session environment details (package versions, Python version, OS) as plain text. Useful for reproducing the notebook environment.

```python
import session_info

session_info.show(html=False)

```

--------------------------------

### Run sequential inference with run_sequential

Source: https://sbijax.readthedocs.io/en/latest/notebooks/more_detailed_intro.html

Runs sequential (multi-round) inference with the run_sequential driver. It simulates from the current posterior each round, appends to the dataset, and refits, keeping the estimator single-round and stateless. The example uses 2 rounds with 2,000 simulations per round and then samples from the fitted model.

```python
from sbijax import run_sequential

params, info = run_sequential(
    jr.PRNGKey(1),
    model,
    prior,
    simulator_fn,
    y_obs,
    n_rounds=2,
    n_simulations_per_round=2_000,
)

_ , _ = model.sample(jr.PRNGKey(3), params, y_obs)

```

--------------------------------

### Generate experimental data with get_experimental_data

Source: https://sbijax.readthedocs.io/en/latest/notebooks/eeg_data_example.html

Generates synthetic and experimental data for the EEG example. It creates a synthetic trajectory, extracts experimental data from X_closed and X_opened, simulates 100,000 training samples using the prior and simulator, and returns a dictionary. It also disables x64 precision in JAX.

```python
def get_experimental_data(prior, n_simulations=100_000):
    # synthetic data example
    theta_synthetic = dict(zip(['C', 'mu', 'sigma', 'gain'], [135, 220, 2000, 0]))
    y_synthetic = simulate_jansen_rit(jr.PRNGKey(1), theta_synthetic)

    # experimental data
    y_experimental_closed = X_closed[[4]]
    y_experimental_opened = X_opened[[4]]

    # simulations
    theta_train = prior.sample(seed=jr.PRNGKey(0), sample_shape=(n_simulations,))
    y_train = simulator(jr.PRNGKey(1), theta_train)

    di = dict(
        theta_synthetic=theta_synthetic,
        y_synthetic=y_synthetic,
        y_experimental_closed=y_experimental_closed,
        y_experimental_opened=y_experimental_opened,
        theta_train=theta_train,
        y_train=y_train
    )
    jax.config.update("jax_enable_x64", False)
    return di

eeg_data = get_experimental_data(prior)
```

--------------------------------

### sbijax function-first API: factories, train, sample, simulate

Source: https://sbijax.readthedocs.io/en/latest/design.html

Shows the function-first API: factories like `npe`/`nle` take only the network, `simulate` generates data from the prior, and the free drivers `train` and `sample` accept the optimizer or sampler as keyword arguments. For likelihood methods, the sampler (kernel + prior) is injected at sample time.

```python
import optax
from jax import numpy as jnp, random as jr
from sbijax import npe, nle, train, sample, simulate
from sbijax.mcmc import make_sampler, nuts
from sbijax.nn import make_maf

# a factory takes only the network
obj = npe(make_maf(2))

# the prior is used only to generate data
data = simulate(jr.key(0), prior, simulator, n=10_000)

# `train` is a free driver; the optimizer is injected here
params, info = train(jr.key(1), obj, data, optimizer=optax.adam(3e-4))

# `sample` is a free driver; for amortized posteriors nothing else is needed
samples, sinfo = sample(jr.key(2), obj, params, y_obs)

# for likelihood/ratio methods the sampler (kernel + prior) is injected here
lik = nle(make_maf(2))
params, _ = train(jr.key(1), lik, data, optimizer=optax.adam(3e-4))
samples, _ = sample(
    jr.key(2), lik, params, y_obs, sampler=make_sampler(nuts, prior=prior)
)
```

--------------------------------

### Show session info with session_info

Source: https://sbijax.readthedocs.io/en/latest/notebooks/getting_started.html

Displays the session environment details using the session_info package. The output shows the versions of key packages (haiku, jax, jaxlib, matplotlib, sbijax, session_info, tensorflow_probability), IPython and Jupyter client versions, Python version, and operating system. Use to reproduce the environment or report dependency versions.

```python
import session_info

session_info.show(html=False)

```

```text
-----
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

```

--------------------------------

### Import JAX, sbijax, and matplotlib

Source: https://sbijax.readthedocs.io/en/latest/notebooks/more_detailed_intro.html

Imports the required libraries for the notebook: JAX, sbijax, and matplotlib for inline plotting.

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

```

--------------------------------

### make_maf: create affine masked autoregressive flow

Source: https://sbijax.readthedocs.io/en/latest/_modules/sbijax/_src/nn/make_flow.html

Public function to create an affine masked autoregressive flow. Validates that n_layers matches the length of n_layer_dimensions when both are provided, otherwise defaults to n_dimension repeated n_layers times. The docstring includes an example call: make_maf(10, n_layer_dimensions=(10, 10, 5, 5, 5)).

```python
def make_maf(
  n_dimension: int,
  n_layers: int | None = 5,
  n_layer_dimensions: Iterable[int] | None = None,
  hidden_sizes: Iterable[int] = (64, 64),
  activation: Callable[..., Any] = jax.nn.tanh,
) -> hk.Transformed:
  """Create an affine (surjective) masked autoregressive flow.

  The MAFs use `n_layers` layers and are parameterized using MADE networks
  with `hidden_sizes` neurons per layer. For each dimensionality reducing
  layer, a conditional Gaussian density is used that uses the same number of
  layer and nodes per layers as `hidden_sizes`. The argument
  `n_layer_dimensions` determines which layer is dimensionality-preserving
  or -reducing. For example, for `n_layer_dimensions=(5, 5, 3, 3)` and
  `n_dimension=5`, the third layer would reduce the dimensionality by two
  and use a surjection layer. THe other layers are dimensionality-preserving.

  Args:
      n_dimension: a list of integers that determine the dimensionality
          of each flow layer
      n_layers: number of layers
      n_layer_dimensions: list of integers that determine if a layer is
          dimensionality-preserving or -reducing
      hidden_sizes: sizes of hidden layers for each normalizing flow
      activation: a jax activation function

  Examples:
      >>> neural_network = make_maf(10, n_layer_dimensions=(10, 10, 5, 5, 5))

  Returns:
      a (surjective) normalizing flow model
  """
  if isinstance(n_layers, int) and n_layer_dimensions is not None:
    assert n_layers == len(list(n_layer_dimensions))
  elif isinstance(n_layers, int):
    n_layer_dimensions = [n_dimension] * n_layers

  return _make_maf(
    n_dimension=n_dimension,
    n_layer_dimensions=n_layer_dimensions,
    hidden_sizes=hidden_sizes,
    activation=activation,
  )
```

--------------------------------

### make_sampler with nuts kernel

Source: https://sbijax.readthedocs.io/en/latest/api/sbijax.mcmc.html

Builds a posterior sampler using the nuts kernel and a prior, then samples from it. The sampler is passed to sbijax.sample().

```python
from sbijax.mcmc import make_sampler, nuts

sampler = make_sampler(nuts, prior=prior)
samples, info = sample(key, estimator, params, y_obs, sampler=sampler)
```

--------------------------------

### SNLE posterior estimation with sbijax

Source: https://sbijax.readthedocs.io/en/latest/notebooks/examples.html

Use surjective neural likelihood estimation (SNLE) to compute a posterior distribution. The workflow starts by importing the SNLE model and a masked autoregressive flow (MAF), then defines the data dimension, network architecture, and prior. After running sequential rounds of simulations, the model samples from the posterior and the results are plotted.

```python
from sbijax import snle, run_sequential
from sbijax.nn import make_maf

```

```python
n_dim_data = 8
n_layer_dimensions, hidden_sizes = (8, 8, 5, 5, 5), (64, 64)
neural_network = make_maf(
    n_dim_data,
    n_layer_dimensions=n_layer_dimensions,
    hidden_sizes=hidden_sizes
)

prior = prior_fn()
model_snle = snle(prior, neural_network)

```

```python
snle_params, info = run_sequential(
    jr.PRNGKey(1),
    model_snle,
    prior,
    simulator_fn,
    y_obs,
    n_rounds=15,
    n_simulations_per_round=1_000,
)

```

```python
snle_samples, _ = model_snle.sample(
    jr.PRNGKey(5), snle_params, y_obs, n_samples=5_000, n_warmup=2_500, n_chains=10
)
  

```

```python
plot_posteriors(
    snle_samples["theta"].reshape(-1, 5),
)
plt.tight_layout()
plt.show()
  

```

--------------------------------

### make_spf and _make_spf: build a rational-quadratic spline coupling flow

Source: https://sbijax.readthedocs.io/en/latest/_modules/sbijax/_src/nn/make_flow.html

The public make_spf function validates layer counts and delegates to the private _make_spf, which builds a rational-quadratic spline coupling flow using distrax, tensorflow_probability, and haiku. The private function constructs MaskedCoupling layers for dimensionality-preserving steps and MaskedCouplingInferenceFunnel layers for reducing steps, then wraps the chain in a TransformedDistribution. The example call make_spf(10, -1.0, 1.0, n_layer_dimensions=(10, 10, 5, 5, 5)) creates a flow with 10 dimensions and a spline range from -1.0 to 1.0.

```python
def make_spf(
  n_dimension: int,
  range_min: float,
  range_max: float,
  n_layers: int | None = 5,
  n_layer_dimensions: Iterable[int] | None = None,
  hidden_sizes: Iterable[int] = (64, 64),
  n_params: int = 10,
  activation: Callable[..., Any] = jax.nn.tanh,
) -> hk.Transformed:
  """Create a rational-quadratic (surjective) spline coupling flow.

  The MAFs use `n_layers` layers and are parameterized using MADE networks
  with `hidden_sizes` neurons per layer. For each dimensionality reducing
  layer, a conditional Gaussian density is used that uses the same number of
  layer and nodes per layers as `hidden_sizes`. The argument
  `n_layer_dimensions` determines which layer is dimensionality-preserving
  or -reducing. For example, for `n_layer_dimensions=(5, 5, 3, 3)` and
  `n_dimension=5`, the third layer would reduce the dimensionality by two
  and use a surjection layer. THe other layers are dimensionality-preserving.

  Args:
      n_dimension: a list of integers that determine the dimensionality
          of each flow layer
      range_min: minimum range on which the spline is defined
      range_max: maximum range on which the spline is defined
      n_layers: number of layers
      n_layer_dimensions: list of integers that determine if a layer is
          dimensionality-preserving or -reducing
      hidden_sizes: sizes of hidden layers for each normalizing flow
      n_params: number of parameters of each spline
      activation: a jax activation function

  Examples:
      >>> neural_network = make_spf(10, -1.0, 1.0, n_layer_dimensions=(10, 10, 5, 5, 5))

  Returns:
      a (surjective) normalizing flow model
  """
  if isinstance(n_layers, int) and n_layer_dimensions is not None:
    assert n_layers == len(list(n_layer_dimensions))
  if isinstance(n_layers, int):
    n_layer_dimensions = [n_dimension] * n_layers

  return _make_spf(
    n_dimension=n_dimension,
    range_min=range_min,
    range_max=range_max,
    n_layer_dimensions=n_layer_dimensions,
    hidden_sizes=hidden_sizes,
    n_params=n_params,
    activation=activation,
  )
```

```python
def _make_spf(
  n_dimension,
  n_layer_dimensions,
  range_min,
  range_max,
  n_params,
  hidden_sizes,
  activation,
) -> hk.Transformed:
  def _bijector_fn(params):
    return distrax.RationalQuadraticSpline(
      params, range_min=range_min, range_max=range_max
    )

  def _decoder_fn(dims):
    def fn(z):
      params = surjectors_mlp(dims, activation=activation)(z)
      mu, log_scale = jnp.split(params, 2, -1)
      return tfd.Independent(tfd.Normal(mu, jnp.exp(log_scale)))

    return fn

  def _conditioner(n_dim):
    return hk.Sequential(
      [
        surjectors_mlp(
          list(hidden_sizes) + [n_params * n_dim],
          activation=activation,
        ),
        hk.Reshape((n_dimension, n_params)),
      ]
    )

  @hk.transform
  def _flow(method, **kwargs):
    layers = []
    curr_dim = n_dimension
    for i, n_dim_curr_layer in enumerate(n_layer_dimensions):
      # layer is dimensionality preserving
      if n_dim_curr_layer == curr_dim:
        layer = MaskedCoupling(
          mask=make_alternating_binary_mask(curr_dim, i % 2 == 0),
          conditioner=_conditioner(curr_dim),
          bijector_fn=_bijector_fn,
        )
      # layer is dimensionality reducing
      elif n_dim_curr_layer < curr_dim:
        n_latent = n_dim_curr_layer
        layer = MaskedCouplingInferenceFunnel(
          n_keep=n_latent,
          decoder=_decoder_fn(list(hidden_sizes) + [2 * (curr_dim - n_latent)]),
          conditioner=surjectors_mlp(
            list(hidden_sizes) + [2 * curr_dim],
            activation=activation,
          ),
          bijector_fn=_bijector_fn,
        )
        curr_dim = n_latent
      else:
        raise ValueError(
          f"n_dimension at layer {i} is layer than the dimension of"
          f" the following layer {i + 1}"
        )
      layers.append(layer)
    chain = Chain(layers)

    base_distribution = tfd.Independent(
      tfd.Normal(jnp.zeros(n_dimension), jnp.ones(n_dimension)),
      1,
    )
    td = TransformedDistribution(base_distribution, chain)
    return td(method, **kwargs)

  return _flow
```

--------------------------------

### Simulate data with prior and simulator function

Source: https://sbijax.readthedocs.io/en/latest/notebooks/getting_started.html

Simulates a dataset of 10,000 samples using the provided prior and simulator function. The resulting data is a dictionary with 'y' and 'theta' keys, containing arrays of the specified shapes.

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

--------------------------------

### NUTS Sampler Implementation

Source: https://sbijax.readthedocs.io/en/latest/_modules/sbijax/_src/mcmc/nuts.html

Internal initialization and kernel definition for the NUTS sampler using BlackJAX.

```python
import blackjax as bj
import jax
from jax import random as jr

from sbijax._src.mcmc.sampler import Kernel
from sbijax._src.mcmc.util import run_blackjax


def _nuts_init(rng_key, initial_positions, lp):
  n_chains = jax.tree_util.tree_leaves(initial_positions)[0].shape[0]
  init_keys = jr.split(rng_key, n_chains)
  warmup = bj.window_adaptation(bj.nuts, lp)
  initial_states, kernel_params = jax.vmap(
    lambda seed, param: warmup.run(seed, param)[0]
  )(init_keys, initial_positions)
  kernel_params = {k: v[0] for k, v in kernel_params.items()}
  _, kernel = bj.nuts(lp, **kernel_params)
  return initial_states, kernel


nuts = Kernel(init_fn=_nuts_init)
```

--------------------------------

### simulate(rng_key, prior, simulator, *, proposal=None, n)

Source: https://sbijax.readthedocs.io/en/latest/_modules/sbijax/_src/simulate/simulate.html

Draws a dataset of parameters and observations for simulation-based inference. Parameters are drawn from the proposal if given, otherwise from the prior, and observations are produced by running the simulator on them.

```APIDOC
## simulate(rng_key, prior, simulator, *, proposal=None, n)

### Description
Draws a dataset of parameters and observations for simulation-based inference. Parameters are drawn from the proposal if given, otherwise from the prior, and observations are produced by running the simulator on them.

### Parameters
- **rng_key** (jax random key) - Required - A jax random key.
- **prior** (distribution) - Required - A distribution to draw parameters from when proposal is None.
- **simulator** (callable) - Required - A callable (rng_key, theta) -> y.
- **proposal** (callable, optional) - Optional - An optional callable (rng_key, n) -> theta used to draw parameters instead of the prior (e.g. a fitted posterior in a sequential round).
- **n** (int) - Required - The number of parameter/observation pairs to draw.

### Returns
A dictionary with keys y and theta.
```
Tokens:
Raw