BlackJAX Integration

Internal API reference for the BlackJAX backend implementation used by MCMCInference.

BlackJAX implementations for filter-based posterior inference.

_run_chain_scan(rng_key, make_step, initial_state, num_steps, info_fn=None)

Scan num_steps MCMC steps, passing a fresh density key to each.

_run_chains(chain_keys, run_chain, initial_states)

Apply one BlackJAX chain runner with standard chain semantics.

init_model(rng_key: jnp.ndarray, model: Callable, *, model_args: tuple, model_kwargs: dict, init_strategy=init_to_median)

Like numpyro's initialize_model but returns a key-aware potential function.

NumPyro's initialize_model fixes the seed when building the potential function, causing Common Random Numbers (CRNs): stochastic model components (particle filters, EnKFs) see the same random seed at every MCMC step.

This function instead returns a potential_fn_gen whose potential functions accept an explicit density_key, so a fresh key can be passed at each step.

Returns:

Type	Description
	(init_params, potential_fn_gen, postprocess_fn) where
	potential_fn_gen(*args) returns potential_fn(position, density_key).

run_blackjax_mcmc(mcmc_config: BaseMCMCConfig, rng_key: jnp.ndarray, model: Callable, obs_times: jnp.ndarray, obs_values: jnp.ndarray, ctrl_times: jnp.ndarray | None = None, ctrl_values: jnp.ndarray | None = None, *model_args, **model_kwargs) -> dict

Run BlackJAX-based inference and return posterior samples.

run_blackjax_mcmc_with_diagnostics(mcmc_config: BaseMCMCConfig, rng_key: jnp.ndarray, model: Callable, obs_times: jnp.ndarray, obs_values: jnp.ndarray, ctrl_times: jnp.ndarray | None = None, ctrl_values: jnp.ndarray | None = None, *model_args, **model_kwargs) -> tuple[dict, dict[str, jax.Array]]

Run BlackJAX inference and return samples plus compact diagnostics.