Converting BlackJAX output to DataTree

Converting BlackJAX output to DataTree#

DataTree is the data format ArviZ relies on.

This page covers multiple ways to generate a DataTree from BlackJAX sampling output.

Log-probability function#

BlackJAX requires an explicit log-probability function. We use a non-centred parameterisation: theta = mu + exp(log_tau) * theta_tilde. The position is a plain dict with keys mu, log_tau, and theta_tilde.

def log_prob(position):
    """Log-joint of the non-centred eight-schools model."""
    mu = position["mu"]
    log_tau = position["log_tau"]
    tau = jnp.exp(log_tau)
    theta_tilde = position["theta_tilde"]
    theta = mu + tau * theta_tilde

    # priors
    lp = -0.5 * (mu / 5.0) ** 2                                # Normal(0, 5)
    lp += log_tau - jnp.log1p((log_tau / jnp.log(5.0)) ** 2)  # Half-Cauchy(0, 5) on tau
    lp += -0.5 * jnp.sum(theta_tilde ** 2)                    # Normal(0, 1) per school

    # likelihood
    lp += -0.5 * jnp.sum(((y_obs - theta) / sigma) ** 2)
    return lp


init_position = {
    "mu": 0.0,
    "log_tau": 0.0,
    "theta_tilde": jnp.zeros(J),
}

Single-chain NUTS#

The standard BlackJAX pattern: use window_adaptation to tune the step size and mass matrix, then run NUTS via jax.lax.scan. The sampler returns a (state, info) tuple at every step; state.position contains the parameter dict and info contains diagnostics such as acceptance_rate and is_divergent.

num_warmup = 500
num_samples = 1000
rng_key = jax.random.PRNGKey(0)

# --- warmup / adaptation ---
warmup = blackjax.window_adaptation(blackjax.nuts, log_prob)
rng_key, warmup_key = jax.random.split(rng_key)
(state, parameters), _ = warmup.run(warmup_key, init_position, num_steps=num_warmup)

# --- sampling ---
kernel = blackjax.nuts(log_prob, **parameters)

def one_step(state, rng_key):
    state, info = kernel.step(rng_key, state)
    return state, (state, info)

rng_key, sample_key = jax.random.split(rng_key)
draw_keys = jax.random.split(sample_key, num_samples)
_, (states, infos) = jax.lax.scan(one_step, state, draw_keys)

Basic conversion#

Pass states.position as posterior and infos as info. Single-chain output has shape (n_draws, *event_shape); from_blackjax inserts a chain dimension of size 1 automatically when sample_dims includes "chain" (the default).

idata = az.from_blackjax(
    posterior=states.position,
    info=infos,
    observed_data={"y": np.array(y_obs)},
    coords={"school": school_names},
    dims={"theta_tilde": ["school"]},
)
idata

The posterior group has dimensions (chain, draw). The sample_stats group contains standard ArviZ names (acceptance_rate, diverging, energy, n_steps, step_size, tree_depth) mapped from the BlackJAX info fields.

print(idata.posterior.dims)
print(list(idata.sample_stats.data_vars))

Adding `reached_max_tree_depth`#

Pass the tree depth cap you used during sampling to get a boolean flag in sample_stats that indicates when the sampler hit the limit.

idata_mtd = az.from_blackjax(
    posterior=states.position,
    info=infos,
    max_tree_depth=10,
)
print("reached_max_tree_depth" in idata_mtd.sample_stats)
print(idata_mtd.sample_stats["reached_max_tree_depth"].values[0])

Multi-chain NUTS via `jax.vmap`#

BlackJAX has no built-in chain management. The idiomatic approach is to vmap a function that runs a single chain over a batch of random keys. The resulting arrays have shape (n_chains, n_draws, *event_shape). Pass num_chains so from_blackjax can identify the leading chain axis.

num_chains = 4
rng_key = jax.random.PRNGKey(1)
chain_keys = jax.random.split(rng_key, num_chains)


def run_chain(key):
    warmup = blackjax.window_adaptation(blackjax.nuts, log_prob)
    (state, parameters), _ = warmup.run(key, init_position, num_steps=num_warmup)
    kernel = blackjax.nuts(log_prob, **parameters)

    def one_step(state, rng_key):
        state, info = kernel.step(rng_key, state)
        return state, (state, info)

    draw_keys = jax.random.split(jax.random.fold_in(key, 1), num_samples)
    _, (states, infos) = jax.lax.scan(one_step, state, draw_keys)
    return states, infos


mc_states, mc_infos = jax.vmap(run_chain)(chain_keys)

idata_mc = az.from_blackjax(
    posterior=mc_states.position,
    info=mc_infos,
    num_chains=num_chains,
    observed_data={"y": np.array(y_obs)},
    coords={"school": school_names},
    dims={"theta_tilde": ["school"]},
)
idata_mc

print(idata_mc.posterior.sizes)

Prior and prior predictive#

BlackJAX has no prior-sampling utility, so prior samples must be generated separately (e.g. with NumPy or JAX). Pass them as a dict to prior.

Variables whose names appear in posterior are placed in the prior group; any remaining variables (not sampled during NUTS) are placed in prior_predictive.

rng = np.random.default_rng(42)
n_prior = 1000

prior_mu = rng.normal(0, 5, size=(1, n_prior))
prior_log_tau = rng.standard_cauchy(size=(1, n_prior))
prior_tau = np.exp(prior_log_tau)
prior_theta_tilde = rng.standard_normal(size=(1, n_prior, J))
prior_theta = prior_mu[..., None] + prior_tau[..., None] * prior_theta_tilde

# y_hat is NOT a parameter in the posterior -> goes to prior_predictive
prior_y_hat = rng.normal(prior_theta, np.array(sigma))

prior_dict = {
    "mu": prior_mu,
    "log_tau": prior_log_tau,
    "theta_tilde": prior_theta_tilde,
    "y_hat": prior_y_hat,
}

idata_prior = az.from_blackjax(
    posterior=states.position,
    info=infos,
    prior=prior_dict,
    observed_data={"y": np.array(y_obs)},
    coords={"school": school_names},
    dims={"theta_tilde": ["school"], "y_hat": ["school"]},
)
idata_prior

# mu, log_tau, theta_tilde -> prior  (overlap with posterior)
# y_hat                    -> prior_predictive  (not in posterior)
print("prior:", list(idata_prior.prior.data_vars))
print("prior_predictive:", list(idata_prior.prior_predictive.data_vars))

`sample_dims` and LOO compatibility#

By default from_blackjax uses sample_dims=["chain", "draw"], matching the ArviZ convention and making the output immediately compatible with functions like az.loo and az.rhat that require a chain dimension.

If you prefer a flat sample representation — e.g. after thinning many chains into one array — pass sample_dims=["sample"]. No chain dimension will be inserted.

idata_flat = az.from_blackjax(
    posterior=states.position,
    info=infos,
    sample_dims=["sample"],
)
print(idata_flat.posterior.dims)
print(idata_flat.posterior.sizes)

If you later need chain and draw dimensions (for example to call az.loo), you can add them with map_over_datasets and expand_dims:

from arviz_base import map_over_datasets

idata_loo_ready = map_over_datasets(
    lambda ds: ds.expand_dims(chain=1, axis=0).rename({"sample": "draw"}),
    idata_flat,
)
print(idata_loo_ready.posterior.dims)

NamedTuple and bare-array positions#

from_blackjax accepts any of the three position formats BlackJAX can produce:

dict (most common) — keys become variable names.
NamedTuple — field names become variable names.
bare array — stored under the variable name "x".

from collections import namedtuple

Position = namedtuple("Position", ["mu", "log_tau", "theta_tilde"])
nt_position = Position(
    mu=states.position["mu"],
    log_tau=states.position["log_tau"],
    theta_tilde=states.position["theta_tilde"],
)

idata_nt = az.from_blackjax(posterior=nt_position)
print(list(idata_nt.posterior.data_vars))  # ['mu', 'log_tau', 'theta_tilde']

# Bare array -- stored under 'x'
idata_bare = az.from_blackjax(posterior=states.position["mu"])
print(list(idata_bare.posterior.data_vars))  # ['x']

Watermark#

%load_ext watermark
%watermark -n -u -v -iv -w

Converting BlackJAX output to DataTree

Contents