sbijax

The top-level module, sbijax, contains all implemented methods for neural simulation-based inference and approximate Bayesian inference as well as diagnostics and other utility.

Every method is a factory function that takes only the network and returns a record of pure functions, following the low-level functional idiom of dm-haiku and blackjax. Training and sampling are free driver functions (train(), sample()); the optimizer is injected at train and, for likelihood/ratio methods, the sampler (which carries the prior) at sample:

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)
)

See Design philosophy for the full design and Migration guide: 0.3 → 0.4 for moving from the class-based API.

npe(network, *[, num_atoms])	Construct a neural posterior estimator.
fmpe(network)	Construct a flow-matching posterior objective.
npse(network)	Construct a neural posterior score estimator.
nle(network)	Construct a neural likelihood objective.
snle(network)	Construct a surjective neural likelihood objective.
nre(network, *[, num_classes, gamma])	Construct a neural ratio objective.
sabc(prior, simulator, *[, summary_fn, ...])	Construct a simulated annealing ABC sampler.
smcabc(prior, simulator, summary_fn, distance_fn)	Construct a sequential Monte Carlo ABC sampler.
nass(network)	Construct a neural approximate sufficient statistics summary network.
nasss(network)	Construct a neural approximate slice sufficient statistics summary network.
summarized_estimator(estimator, summary_net, ...)	Adapt an objective to operate on learned summaries.
train(rng_key, objective, data, *[, ...])	Train any objective's TrainFns with early stopping.
sample(rng_key, objective, params, observable, *)	Draw posterior samples from a trained objective.
run_sequential(rng_key, objective, prior, ...)	Run multi-round sequential inference.
simulate(rng_key, prior, simulator, *[, ...])	Draw a dataset of parameters and observations.
stack(data, new_data)	Append two datasets along the sample axis.
sbc(rng_key, objective, params, prior, ...)	Compute simulation-based calibration ranks for a fitted objective.
rhat(samples)	Split-R-hat for each parameter dimension.
ess(samples)	Effective sample size for each parameter dimension.

Data pipeline

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

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

sbijax.stack(data, new_data)

Append two datasets along the sample axis.

Parameters:

• data – a dataset as returned by simulate()

• new_data – a second dataset with the same structure

Returns:

a dataset whose leaves are the concatenation of both inputs

Posterior estimation

sbijax.npe(network, *, num_atoms=10)

Construct a neural posterior estimator.

In round 0 use the returned ObjectiveFns directly for amortized maximum-likelihood training. For round > 0 (i.e. when training on data simulated from a fitted posterior rather than the prior), call obj.extra(prior) to obtain the atomic proposal-posterior objective of Greenberg et al. [2019].

Parameters:

• network – a conditional density estimator with log_prob and sample methods (e.g. from make_maf())

• num_atoms – the number of atoms in the contrastive proposal-posterior loss used by extra(prior)

Returns:

an ObjectiveFns; its extra field is a callable (prior) -> ObjectiveFns for sequential rounds

sbijax.fmpe(network)

Construct a flow-matching posterior objective.

Parameters:

network – a continuous normalizing flow with loss and sample methods

Returns:

an ObjectiveFns

sbijax.npse(network)

Construct a neural posterior score estimator.

Parameters:

network – a score network with loss, sample and log_prob methods (e.g. sbijax.experimental.nn.make_score_model())

Returns:

an ObjectiveFns

Likelihood estimation

sbijax.nle(network)

Construct a neural likelihood objective.

Parameters:

network – a conditional density estimator exposing a log_prob method

Returns:

an ObjectiveFns

sbijax.snle(network)

Construct a surjective neural likelihood objective.

Parameters:

network – a surjective conditional density estimator with a log_prob method that reduces the dimensionality of the data

Returns:

an ObjectiveFns

Likelihood-ratio estimation

sbijax.nre(network, *, num_classes=10, gamma=1.0)

Construct a neural ratio objective.

Parameters:

• network – a classifier network mapping concat(y, theta) to logits

• num_classes – number of contrastive classes

• gamma – relative weight of the contrastive classes

Returns:

an ObjectiveFns

Approximate Bayesian computation

sbijax.sabc(prior, simulator, *, summary_fn=<function <lambda>>, distance_fn=<function abs_distance>)

Construct a simulated annealing ABC sampler.

Parameters:

• prior – a tfd distribution serving as the prior over parameters

• simulator – a callable (rng_key, theta) -> y

• summary_fn – maps simulated data to summary statistics

• distance_fn – distance between simulated and observed summaries

Returns:

an ABCSampler

sbijax.smcabc(prior, simulator, summary_fn, distance_fn)

Construct a sequential Monte Carlo ABC sampler.

Parameters:

• prior – a tfd distribution serving as the prior over parameters

• simulator – a callable (rng_key, theta) -> y

• summary_fn – maps simulated data to summary statistics

• distance_fn – distance between simulated and observed summaries, one value per simulation, of shape (n,) or (n, 1) (e.g. l2_distance())

Returns:

an ABCSampler

Summary statistics

sbijax.nass(network)

Construct a neural approximate sufficient statistics summary network.

Parameters:

network – a NASS summary network with forward, summary and critic methods

Returns:

a SummaryFns

sbijax.nasss(network)

Construct a neural approximate slice sufficient statistics summary network.

Parameters:

network – a NASSS summary network with forward, summary, secondary_summary and critic methods

Returns:

a SummaryFns

A summary network is chained into a downstream estimator with:

sbijax.summarized_estimator(estimator, summary_net, summary_params)

Adapt an objective to operate on learned summaries.

Given a pre-fitted summary network, returns an ObjectiveFns whose training summarizes each batch before delegating to the wrapped estimator and whose sample_fn summarizes the observation before sampling. Because it returns an ObjectiveFns record it stays conformant and is driven by the generic train() / sampling helpers.

Fit the summary network first, then wrap the estimator:

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)

Parameters:

• estimator – the downstream ObjectiveFns consuming the summaries

• summary_net – a fitted SummaryFns

• summary_params – the summary network’s fitted parameters

Returns:

an ObjectiveFns

Training and sampling

Trainable objectives are trained and sampled with the two free drivers. The sampler for likelihood/ratio methods is built with sbijax.mcmc.make_sampler().

sbijax.train(rng_key, objective, data, *, optimizer=None, info=None, n_iter=1000, batch_size=100, percentage_data_as_validation_set=0.1, n_early_stopping_patience=10, n_early_stopping_delta=0.001)

Train any objective’s TrainFns with early stopping.

Parameters:

• rng_key – a jax random key

• objective – an ObjectiveFns/SummaryFns (anything with train)

• data – a {"y", "theta"} dataset pytree

• optimizer – an optax optimizer, bound into the primitives here

• info – previous round’s Info (None for round 0)

• n_iter – number of epochs

• batch_size – minibatch size

• percentage_data_as_validation_set – validation split fraction

• n_early_stopping_patience – early-stopping patience

• n_early_stopping_delta – minimum early-stopping improvement

Returns:

a tuple (params, Info)

sbijax.sample(rng_key, objective, params, observable, *, sampler=None, prior=None, **kwargs)

Draw posterior samples from a trained objective.

A thin dispatch to objective.sample_fn kept for symmetry with sbijax.train().

Parameters:

• rng_key – a jax random key

• objective – an ObjectiveFns

• params – the trained parameters

• observable – the observation to condition on

• sampler – a sampler from make_sampler() (required for MCMC methods, ignored by amortized methods)

• prior – the prior of the model. The amortized estimators reject draws outside its support and return the flattened parameter vector under a single "theta" key; passing the prior here also reshapes them into the prior’s pytree, matching what the MCMC and ABC methods return. Without it no draw is rejected and the flat layout is preserved.

• **kwargs – forwarded to sample_fn

Returns:

(samples, info)

Sequential inference

sbijax.run_sequential(rng_key, objective, prior, simulator, observable, *, n_rounds, n_simulations_per_round, sampler=None, proposal_fn=None, **train_kwargs)

Run multi-round sequential inference.

Round 0 simulates from the prior; later rounds simulate from a proposal built from the current posterior, append, and refit. NPE switches to its atomic objective (objective.extra(prior)) in rounds > 0; proposal-invariant methods reuse objective.

Parameters:

• rng_key – a jax random key

• objective – an ObjectiveFns

• prior – the prior distribution

• simulator – (rng_key, theta) -> y

• observable – the observation to condition on

• n_rounds – number of simulate/append/refit rounds

• n_simulations_per_round – pairs drawn each round

• sampler – a sampler (from make_sampler) for MCMC-based objectives

• proposal_fn – optional (objective, params, observable, sampler) -> ((rng, n) -> theta); defaults to sampling the fitted posterior

• **train_kwargs – forwarded to train each round

Returns:

(params, Info) from the final round

Diagnostics

sbijax.sbc(rng_key, objective, params, prior, simulator, *, sampler=None, n_calibrations=100, n_samples=1000, **sample_kwargs)

Compute simulation-based calibration ranks for a fitted objective.

For each of n_calibrations draws theta* ~ prior and y* ~ simulator(theta*), draws n_samples from the posterior given y* and ranks each dimension of theta* among them. Calibrated posteriors yield ranks uniform on [0, n_samples].

Parameters:

• rng_key – a jax random key

• objective – an ObjectiveFns returned by a factory such as npe()

• params – the fitted parameters

• prior – the prior distribution

• simulator – a callable (rng_key, theta) -> y

• sampler – a sampler from make_sampler(); required for MCMC methods, ignored by amortized methods

• n_calibrations – number of calibration draws

• n_samples – posterior draws per calibration draw

• **sample_kwargs – forwarded to sample

Returns:

an integer array of shape (n_calibrations, n_dims) of ranks

sbijax.rhat(samples)

Split-R-hat for each parameter dimension.

Parameters:

samples – a named pytree with leaves of shape (n_chains, n_draws, dim)

Returns:

a pytree of the same structure with per-dimension R-hat values

sbijax.ess(samples)

Effective sample size for each parameter dimension.

Parameters:

samples – a named pytree with leaves of shape (n_chains, n_draws, dim)

Returns:

a pytree of the same structure with per-dimension ESS values