blackjax.diagnostics

MCMC diagnostics.

Classes

DivergenceConcentrationReport

Per-run divergence-concentration diagnostic report.

Functions

potential_scale_reduction(→ blackjax.types.Array)	Gelman and Rubin (1992)'s potential scale reduction for computing multiple MCMC chain convergence.
rhat(→ blackjax.types.Array)	Rank-normalized split-R̂ (Vehtari et al. 2021).
effective_sample_size(→ blackjax.types.Array)	Compute estimate of the effective sample size (ess).
ess_bulk(→ blackjax.types.Array)	Bulk effective sample size (rank-normalized split-chain ESS).
ess_tail(→ blackjax.types.Array)	Tail effective sample size.
pareto_khat(→ blackjax.types.Array)	Pareto shape parameter k̂ for tail diagnosis.
psis_weights(→ tuple[blackjax.types.Array, ...)	Pareto Smoothed Importance Sampling (PSIS) log weights.
divergence_concentration(→ DivergenceConcentrationReport)	Flag a run where sampling-phase divergences concentrate on a minority of chains.
divergence_concentration_from_counts(...)	Same as divergence_concentration(), from precomputed per-chain counts.
format_divergence_warning(→ str)	Render a DivergenceConcentrationReport as a human-readable message.

Module Contents

potential_scale_reduction(input_array: blackjax.types.ArrayLike, chain_axis: int = 0, sample_axis: int = 1) → blackjax.types.Array

Gelman and Rubin (1992)’s potential scale reduction for computing multiple MCMC chain convergence.

Parameters:

• input_array – An array representing multiple chains of MCMC samples. The array must contains a chain dimension and a sample dimension.

• chain_axis – The axis indicating the multiple chains. Default to 0.

• sample_axis – The axis indicating a single chain of MCMC samples. Default to 1.

Return type:

NDArray of the resulting statistics (r-hat), with the chain and sample dimensions squeezed.

Notes

The diagnostic is computed by:

\[\hat{R} = \sqrt{\frac{\hat{V}}{W}}\]

where \(W\) is the within-chain variance and \(\hat{V}\) is the posterior variance estimate for the pooled traces. This is the potential scale reduction factor, which converges to unity when each of the traces is a sample from the target posterior. Values greater than one indicate that one or more chains have not yet converged [stac, GR92].

rhat(input_array: blackjax.types.ArrayLike, chain_axis: int = 0, sample_axis: int = 1) → blackjax.types.Array

Rank-normalized split-R̂ (Vehtari et al. 2021).

The modern improved R̂ diagnostic. Combines two split-chain R̂ values — one on rank-normalized draws and one on rank-normalized folded draws — and returns the maximum. The folded component catches scale/variance non-convergence that the bulk component can miss.

This matches the default az.rhat(method="rank") convention in ArviZ.

Parameters:

• input_array – An array representing multiple chains of MCMC samples. The array must contain a chain dimension and a sample dimension. At least 2 chains and at least 4 draws per chain are required.

• chain_axis – The axis indicating the multiple chains. Default 0.

• sample_axis – The axis indicating a single chain of MCMC samples. Default 1.

Returns:

• NDArray of the resulting R̂ values, with chain and sample dimensions

• squeezed. Values close to 1.0 indicate convergence; values above 1.01

• suggest chains have not converged.

Notes

Algorithm (Vehtari et al. 2021, § 4):

1. Split each chain in half → 2× chains.

2. Rank-normalize with the Blom plotting position \(z_r = \Phi^{-1}((r - 3/8) / (n + 1/4))\) over the joint pool.

3. Compute the standard split-R̂ on the rank-normalized draws (bulk).

4. Compute the folded draws \(|x - \mathrm{median}(x)|\), rank-normalize them, and compute split-R̂ again (tail).

5. Return \(\max(\hat{R}_{\text{bulk}}, \hat{R}_{\text{tail}})\).

Warning

NaN from this function does not uniquely mean “the draws contained a missing observation”. The folded component subtracts the pooled median, so when at least half of a component’s pooled draws are +inf (or -inf) the median is infinite and the fold evaluates inf - inf, which is NaN. Such a component returns NaN from rhat() while containing no NaN at all — and while ess_bulk() and ess_tail() still return finite values for it. Infinity handling is a separate open question from the missing-data contract; this is documented, not yet decided.

References

effective_sample_size(input_array: blackjax.types.ArrayLike, chain_axis: int = 0, sample_axis: int = 1) → blackjax.types.Array

Compute estimate of the effective sample size (ess).

Parameters:

• input_array – An array representing multiple chains of MCMC samples. The array must contains a chain dimension and a sample dimension.

• chain_axis – The axis indicating the multiple chains. Default to 0.

• sample_axis – The axis indicating a single chain of MCMC samples. Default to 1.

Returns:

• NDArray of the resulting statistics (ess), with the chain and sample dimensions squeezed.

• Variables whose within-chain variance is numerically zero have an effective

• sample size of zero. Variables containing a NaN draw are undefined and

• report NaN; the reduction is per variable, so an independent finite

• variable is unaffected. Infinities are ordered values and are not treated

• as missing.

Notes

The basic ess (\(N_{\mathit{eff}}\)) diagnostic is computed by:

\[\hat{N}_{\mathit{eff}} = \frac{MN}{\hat{\tau}}\]

\[\hat{\tau} = -1 + 2 \sum_{t'=0}^K \hat{P}_{t'}\]

where \(M\) is the number of chains, \(N\) the number of draws, \(\hat{\rho}_t\) is the estimated _autocorrelation at lag \(t\), and \(K\) is the last integer for which \(\hat{P}_{K} = \hat{\rho}_{2K} + \hat{\rho}_{2K+1}\) is still positive [staa, GCSR95].

The current implementation is similar to Stan, which uses Geyer’s initial monotone sequence criterion [Gey92, Gey11].

ess_bulk(input_array: blackjax.types.ArrayLike, chain_axis: int = 0, sample_axis: int = 1) → blackjax.types.Array

Bulk effective sample size (rank-normalized split-chain ESS).

Computes the bulk ESS from Vehtari et al. (2021): rank-normalizes draws after splitting each chain in half, then applies the standard autocorrelation-based effective_sample_size() estimator. This diagnostic is robust to non-stationarity and multimodality.

Parameters:

• input_array – An array representing multiple chains of MCMC samples. The array must contain a chain dimension and a sample dimension.

• chain_axis – The axis indicating the multiple chains. Default 0.

• sample_axis – The axis indicating a single chain of MCMC samples. Default 1.

Return type:

NDArray of the resulting bulk-ESS, with chain and sample dimensions squeezed.

Notes

Algorithm:

1. Split each chain in half → 2× chains.

2. Pool all draws and rank-normalize with \(z_r = \Phi^{-1}((r-3/8)/(n+1/4))\).

3. Apply effective_sample_size() to the rank-normalized draws.

References

ess_tail(input_array: blackjax.types.ArrayLike, chain_axis: int = 0, sample_axis: int = 1, prob: float = 0.9) → blackjax.types.Array

Tail effective sample size.

Computes the tail ESS from Vehtari et al. (2021) as the minimum of the ESS of the lower- and upper-tail indicator functions applied to split-chain draws.

The tail quantiles are determined by prob: the lower tail uses the (1 - prob) / 2 quantile and the upper tail uses the (1 + prob) / 2 quantile. The default prob=0.90 corresponds to the 5th/95th percentiles, matching ArviZ’s default prob=(0.05, 0.95).

Warning

The agreement with az.ess(method="tail") holds for continuous draws only. The upper-tail indicator here is \(\mathbf{1}(x \ge q_{\text{high}})\), whereas Vehtari et al. and ArviZ use \(\mathbf{1}(x \le q)\) for both tails. On continuous draws the two are exact complements and the ESS is identical, but on tied draws they are not: for iid Bernoulli(0.1), \(P(x \ge q_{95})\) is 0.097 rather than 0.05, and for a 5-level grid it is 0.21. This is a separate known defect in the tail estimator’s tie handling, tracked independently of the rank-normalization fix; note that simply switching to <= does not resolve it, since that indicator is identically 1 on such draws and would be reported as degenerate.

Parameters:

• input_array – An array representing multiple chains of MCMC samples. The array must contain a chain dimension and a sample dimension.

• chain_axis – The axis indicating the multiple chains. Default 0.

• sample_axis – The axis indicating a single chain of MCMC samples. Default 1.

• prob – Central-interval probability that determines the tail quantiles. Lower quantile: (1 - prob) / 2; upper quantile: (1 + prob) / 2. Default 0.90 gives the 5th/95th-percentile tail, matching az.ess(method="tail") (ArviZ default).

Return type:

NDArray of the resulting tail-ESS, with chain and sample dimensions squeezed.

Notes

Algorithm:

1. Split each chain in half → 2× chains.

2. Compute pooled lower/upper quantiles (at (1-prob)/2 and (1+prob)/2) across all split chains and draws.

3. Form indicator series \(\mathbf{1}(x \le q_{\text{low}})\) and \(\mathbf{1}(x \ge q_{\text{high}})\).

4. Compute effective_sample_size() for each indicator.

5. Return \(\min(\text{ESS}_\text{lower}, \text{ESS}_\text{upper})\).

References

pareto_khat(x: blackjax.types.ArrayLike, tail: str = 'both', tail_frac: float = 0.1) → blackjax.types.Array

Pareto shape parameter k̂ for tail diagnosis.

Fits a Generalised Pareto Distribution (GPD) to the upper and/or lower tail of a 1-D sample and returns the estimated shape parameter k̂.

Parameters:

• x – 1-D array of draws (or any array; it is ravelled before use).

• tail – Which tail to fit: "upper", "lower", or "both" (default). When "both", returns the maximum of the two k̂ estimates.

• tail_frac – Fraction of samples used as the tail. Default 0.10 (10 %). A minimum of 5 tail samples is always enforced.

Returns:

• Scalar Array (the Pareto shape estimate k̂. Values below 0.5 indicate)

• reliable tail estimates; 0.5–0.7 are moderate; above 0.7 may be

• unreliable.

Notes

Uses the Zhang & Stephens (2009) empirical-Bayes estimator implemented in the internal _gpdfit(). The upper tail is modelled directly; the lower tail is reflected and modelled as an upper tail.

psis_weights(log_ratios: blackjax.types.Array, r_eff: float = 1.0) → tuple[blackjax.types.Array, blackjax.types.Array]

Pareto Smoothed Importance Sampling (PSIS) log weights.

Implements the PSIS smoothing step from [VGG17]: the M largest importance ratios (in ratio space) are replaced by sorted Generalised Pareto quantiles fitted by the empirical Bayes estimator of Zhang & Stephens (2009), then all weights are normalised.

This is a pure-JAX, JIT-compatible implementation faithful to Algorithm 1 of Vehtari, Gelman & Gabry (2017).

Parameters:

• log_ratios – Log importance ratios log p(θ) − log q(θ), shape (n,). Need not be normalised.

• r_eff – Relative effective sample size of the proposal, S_eff / n. Use the default of 1.0 for i.i.d. draws (e.g. Pathfinder); set to the actual ESS ratio for correlated MCMC chains. Values below 1 increase the tail size M to compensate for correlation.

Returns:

• log_weights – Normalised log importance weights, shape (n,). jnp.exp(log_weights).sum() == 1 up to floating-point precision.

• pareto_k – Pareto shape parameter estimate (scalar Array). Values below 0.5 indicate reliable estimates; 0.5–0.7 are moderate; above 0.7 may give unreliable estimates. jnp.inf means the tail was too small to fit (fewer than 5 samples).

Notes

Tail size: M = min(floor(3*sqrt(n/r_eff)), n//5), matching the paper. The GPD is only applied when k >= 1/3; lighter tails are left unsmoothed (only normalised). Fitting uses empirical Bayes in importance-ratio space, the same approach as ArviZ.

class DivergenceConcentrationReport

Per-run divergence-concentration diagnostic report.

All fields are plain JAX numerics — no strings — so the statistic itself stays JIT-compilable. Pass the report to format_divergence_warning() to render a human-readable message from concrete (non-traced) values.

warn

Bool. Whether a minority of chains (see the module notes on divergence_concentration()) crossed rate_threshold.

flagged

Bool array, shape (n_chains,). Which chains crossed rate_threshold, independent of whether warn ends up true.

num_flagged

Number of flagged chains, flagged.sum().

rates

Per-chain sampling-phase divergence rate, shape (n_chains,).

early_rate, late_rate

Per-chain divergence rate in the first and last quarter of the sampling draws, shape (n_chains,). NaN when computed from divergence_concentration_from_counts() (no per-draw resolution available).

median_other_rate

For each chain, the median rates value across the other n_chains - 1 chains, shape (n_chains,). NaN when n_chains <= 1.

total_divergences

Total divergence count D summed over all chains.

num_chains

Number of chains M.

rate_threshold

The threshold that was applied (echoed back for the message / for callers that only keep the report).

multinomial_p_value

Bonferroni-corrected exchangeable-null tail probability for the single worst chain, min(1, M * P(Binomial(D, 1/M) >= d_max)). Context only; never drives warn. NaN when D=0.

warn: blackjax.types.Array

flagged: blackjax.types.Array

num_flagged: blackjax.types.Array

rates: blackjax.types.Array

early_rate: blackjax.types.Array

late_rate: blackjax.types.Array

median_other_rate: blackjax.types.Array

total_divergences: blackjax.types.Array

num_chains: blackjax.types.Array

rate_threshold: blackjax.types.Array

multinomial_p_value: blackjax.types.Array

divergence_concentration(is_divergent: blackjax.types.ArrayLike, *, rate_threshold: float = 0.02) → DivergenceConcentrationReport

Flag a run where sampling-phase divergences concentrate on a minority of chains.

Parameters:

• is_divergent – Per-draw divergence flags (bool or 0/1) for the sampling (post-warmup) phase only, shape (n_chains, n_draws).

• rate_threshold – Minimum per-chain divergence rate for chain k to count as flagged. Default 0.02 (2%).

Returns:

• DivergenceConcentrationReport. warn is true iff between

• 1 and max(1, n_chains // 4) chains are flagged – a

• minority-outlier trigger; an ensemble where most/all chains cross the

• threshold returns populated fields but no warning.

Notes

• Counting is sampling-phase only; warmup divergences are out of scope.

• Concatenating warmup draws in front of sampling draws dilutes a real signal below threshold rather than raising a false alarm – slice to sampling draws only before calling this.

• multinomial_p_value is context only; it never decides warn.

divergence_concentration_from_counts(chain_divergence_counts: blackjax.types.ArrayLike, n_draws: int, *, rate_threshold: float = 0.02) → DivergenceConcentrationReport

Same as divergence_concentration(), from precomputed per-chain counts.

Use when per-draw flags are not retained but per-chain totals are. early_rate / late_rate are NaN (no per-draw resolution for a quarter profile); format_divergence_warning() omits that part of the message rather than printing “nan%”.

Parameters:

• chain_divergence_counts – Per-chain divergence counts for the sampling phase, shape (n_chains,).

• n_draws – Number of sampling draws per chain.

• rate_threshold – See divergence_concentration().

Return type:

DivergenceConcentrationReport

format_divergence_warning(report: DivergenceConcentrationReport) → str

Render a DivergenceConcentrationReport as a human-readable message.

Returns "" when report.warn is false. Otherwise returns one sentence per flagged chain (newline-separated when more than one chain is flagged). Not JIT-compatible by design — call it on concrete report values (e.g. after a sampling run has completed), not inside traced code.

Parameters:

report – A report produced by divergence_concentration() or divergence_concentration_from_counts().

Return type:

str, empty when there is nothing to warn about.