mixle.inference.target module

Bring-your-own-target, engine-agnostic inference facade.

A small public surface so external consumers can run mixle’s samplers / VI on an arbitrary differentiable target without reaching into mixle.ppl internals — and run it on whichever engine fits their hardware and target:

  • nuts() — multi-chain No-U-Turn Sampler with R-hat / ESS diagnostics, dispatched to a registered backend (numpy / numba / torch / jax) selected by backend="auto". Every backend returns the same NutsResult, so downstream code is backend-blind.

  • nuts_torch() — the torch-native NUTS, kept as a direct entry point (also the "torch" backend).

  • advi() — automatic-differentiation VI over a user batched torch target.

  • available_backends() — the installed inference engines.

  • mixle.inference.diagnostics (re-exported as rhat() / ess()) — convergence diagnostics over plain (n_chains, n_draws, d) arrays.

The backend registry lives in mixle.inference.backends (register, don’t branch); each backend self-registers at import, behind a try/except so a missing optional engine never breaks import mixle.inference.

class NutsResult(samples, chains, rhat, ess, num_target_evals, step_size, extra=<factory>)[source]

Bases: object

Draws and diagnostics from a multi-chain nuts() run.

samples is (chains * draws, d) pooled draws; chains is (n_chains, draws, d).

Parameters:
samples: ndarray
chains: ndarray
rhat: ndarray
ess: ndarray
num_target_evals: int
step_size: float
extra: dict
class AdviResult(samples, mean, scale, objective=None)[source]

Bases: object

Result of advi(): posterior draws plus the fitted variational parameters.

Parameters:
samples: ndarray
mean: ndarray
scale: ndarray
objective: float | None = None
nuts(target, *, backend='auto', dim=None, init=None, num_samples=1000, warmup=1000, chains=1, mass=1.0, target_accept=0.8, max_tree_depth=10, thin=1, rng=None, parallel=None, **backend_kwargs)[source]

No-U-Turn Sampler over an arbitrary log-target, dispatched to a registered engine.

The target contract depends on the backend (the kinds cannot be auto-converted across autodiff systems):

  • numpy / numba: a fused value_and_grad(theta) -> (logp, grad) (njit-jitted for numba). The caller supplies the (analytic) gradient.

  • torch / jax: a scalar logp(theta); the backend builds value_and_grad by autodiff (torch.func / NumPyro).

Parameters:

  • target (Callable[[...], Any]) – the log-target, per the contract above.

  • backend (str) – "auto" (default), or one of available_backends(). "auto" honors an explicit choice elsewhere, otherwise prefers the always-present numpy path for a plain numpy value_and_grad; pass backend="numba" to run an @njit target, "jax" / "torch" for an autodiff scalar logp.

  • dim (int | None) – parameter dimension. Required unless init is given.

  • init (Any) – initial state, shape (dim,) or (chains, dim) for per-chain starts. Defaults to zeros. A 1-D init is reused (jittered after the first) across chains.

  • num_samples (int) – retained post-warmup draws per chain.

  • warmup (int) – step-size adaptation / burn-in iterations per chain.

  • chains (int) – number of independent chains (>= 2 to get a meaningful R-hat).

  • mass (Any) – forwarded to the sampler.

  • target_accept (float) – forwarded to the sampler.

  • max_tree_depth (int) – forwarded to the sampler.

  • thin (int) – forwarded to the sampler.

  • rng (RandomState | int | None) – seed / RandomState for reproducibility.

  • parallel (bool | str | None) – run the chains independent chains concurrently. None/False (default) runs them in the backend’s usual single call; "thread" uses a thread pool (real speedup for the numba/torch backends, whose inner loops release the GIL); True or "process" uses a process pool (real speedup for the pure-numpy backend; requires a picklable target). Ignored for the jax backend, which already vectorizes chains internally. Each chain still gets an independent seed, so results are valid multi-chain draws (R-hat / ESS across chains).

  • **backend_kwargs (Any) – forwarded to the chosen backend (e.g. compile=, device= for torch; chain_method= for jax).

Returns:

NutsResult with pooled samples (chains*draws, d), per-chain chains (chains, draws, d), per-dimension rhat and ess, the total target-evaluation count, the adapted step_size, and extra={"backend": name}.

Return type:

NutsResult

nuts_torch(logp, *, dim=None, init=None, num_samples=1000, warmup=1000, chains=1, mass=1.0, target_accept=0.8, max_tree_depth=10, thin=1, rng=None, compile=True, dtype=None, device=None)[source]

Torch-native NUTS over a torch scalar logp(theta) -> Tensor[()] (GPU / large autodiff targets).

The whole leapfrog/tree runs in torch tensors with a torch.compile``d ``value_and_grad (no numpy round-trip / graph re-trace per gradient). Same multi-chain interface and NutsResult as nuts(). Also registered as the "torch" backend.

Performance note (be deliberate about when to use this): on CPU this is typically slower than the numpy nuts() (per-op torch dispatch + host syncs in the tree dominate when the target is cheap). Its value is GPU (device=) and large autodiff targets. For CPU work prefer the numpy backend, numba (analytic gradient), or the jax/NumPyro backend (XLA, vectorized multi-chain). Chains run in a Python loop (per-chain latency).

Parameters:
Return type:

NutsResult

advi(target_batch, u0, s0, *, samples=1000, mc=16, steps=2000, lr=0.05, batch_size=None, family='meanfield', alpha=1.0, rng=None)[source]

Automatic-differentiation VI over a user batched torch target.

Parameters:

  • target_batch (Callable[[Any], Any]) – target_batch(U: Tensor(B, d)) -> Tensor(B,) returning the (unnormalized) joint log-target for each of B parameter draws. The caller owns any data minibatching/rescaling inside this callable; batch_size here is unused unless the caller wires it in (kept for signature parity with the internal ADVI).

  • u0 – initial variational mean and (marginal) scale in the unconstrained space.

  • s0 – initial variational mean and (marginal) scale in the unconstrained space.

  • samples (int) – number of posterior draws to return.

  • mc (int) – Monte-Carlo samples per step, optimizer steps, Adam learning rate.

  • steps (int) – Monte-Carlo samples per step, optimizer steps, Adam learning rate.

  • lr (float) – Monte-Carlo samples per step, optimizer steps, Adam learning rate.

  • family (str) – 'meanfield' (diagonal) or 'fullrank' (Cholesky covariance).

  • alpha (float) – Renyi/tilted objective exponent (1.0 = standard KL-ELBO).

  • rng (RandomState | int | None) – seed / RandomState.

  • batch_size (int | None)

Returns:

AdviResult with samples (samples, d) drawn from the fitted Gaussian q (in the same space target_batch consumes), plus the fitted mean and scale.

Return type:

AdviResult

rhat(chains)[source]

Gelman-Rubin potential scale reduction factor (R-hat) per parameter dimension.

R-hat compares the variance between chains to the variance within chains. Values near 1.0 indicate the chains have mixed and are sampling a common target; values above ~1.01-1.1 flag non-convergence. Standard multi-chain check (Gelman & Rubin 1992).

Parameters:

chains (Any) – array shaped (n_chains, n_draws) or (n_chains, n_draws, d). Needs at least two chains and two draws.

Returns:

A length-d array of R-hat values (one per parameter).

Return type:

ndarray

ess(samples, max_lag=None)[source]

Effective sample size per parameter from positive autocorrelation lags.

Accepts a single chain (n_draws,) / (n_draws, d) or a stack of chains (n_chains, n_draws, d) (the chains are pooled into one autocorrelation estimate after centering each chain by its own mean). Uses the initial-positive-sequence truncation (Geyer): sum autocorrelations until the first non-positive lag.

Returns:

A length-d array of ESS values.

Parameters:

  • samples (Any)

  • max_lag (int | None)

Return type:

ndarray

split_rhat(chains)[source]

Rank-normalized split-R-hat (Vehtari et al. 2021) – the robust, recommended R-hat.

Splits each chain in half (doubling the chain count, so within-chain non-stationarity is caught), rank-normalizes the pooled draws (robust to heavy tails / non-normality), then applies the Gelman-Rubin R-hat. Convergence is typically declared at < 1.01.

Parameters:

chains (Any)

Return type:

ndarray

folded_split_rhat(chains)[source]

Folded rank-normalized split-R-hat (Vehtari et al. 2021): split-R-hat on |x - median(x)|.

The plain split_rhat() compares chain locations; folding about the median makes it compare chain scales/tails, catching the case where chains share a mean but differ in spread. Use together with split_rhat() (or rhat_max()).

Parameters:

chains (Any)

Return type:

ndarray

rhat_max(chains)[source]

The recommended convergence R-hat: max(split_rhat, folded_split_rhat) (Vehtari et al. 2021).

A single number that flags non-convergence in either the location (bulk) or the scale (folded) of the chains; declare convergence at < 1.01.

Parameters:

chains (Any)

Return type:

ndarray

mcse_mean(chains)[source]

Monte Carlo standard error of the posterior mean: sd(x) / sqrt(ESS) per parameter.

The sampling error in the estimated mean from autocorrelated draws – the posterior standard deviation deflated by the (autocorrelation-based) effective sample size of the raw chains.

Parameters:

chains (Any)

Return type:

ndarray

mcmc_summary(chains)[source]

Per-parameter posterior + convergence summary (the ArviZ-style summary table).

Returns one dict per parameter with the posterior mean/sd and 5/50/95% quantiles, plus the recommended convergence diagnostics: r_hat (= rhat_max(), the max of the bulk and folded rank-normalized split-R-hats), ess_bulk, ess_tail, and mcse_mean (Vehtari et al. 2021). Declare convergence at r_hat < 1.01 with ess_bulk/ess_tail comfortably in the hundreds.

Parameters:

chains (Any)

Return type:

list[dict[str, float]]

ess_bulk(chains)[source]

Bulk effective sample size: ESS of the rank-normalized split chains (mixing in the distribution body).

Parameters:

chains (Any)

Return type:

ndarray

ess_tail(chains, prob=0.05)[source]

Tail effective sample size: the smaller ESS of the lower-/upper-prob tail indicators.

Tail quantiles mix more slowly than the bulk; tail-ESS is min(ESS[1(x <= q_prob)], ESS[1(x >= q_{1-prob})]) over the split chains (Vehtari et al. 2021), surfacing poor tail exploration that bulk-ESS misses.

Parameters:
Return type:

ndarray

geweke_z(chain, first=0.1, last=0.5)[source]

Geweke convergence z-score comparing the start and end of a single chain (per parameter).

Tests stationarity within one chain (Geweke 1992): if the chain has converged, the mean of its first first fraction and its last last fraction agree, so the z-statistic (mean_a - mean_b) / sqrt(se_a^2 + se_b^2) – with autocorrelation-adjusted standard errors (the segment variance divided by its effective sample size) – is approximately standard normal. |z| < 2 indicates convergence; a large |z| flags a chain still drifting. Returns one z per parameter. Complements the multi-chain split_rhat().

Parameters:
Return type:

ndarray

available_backends()[source]

Return the names of registered backends whose engine is importable, in registration order.

Return type:

list[str]

class InferenceBackend(name, available, target_kind, nuts)[source]

Bases: object

A registered inference engine: a name, an availability probe, a target contract, a sampler.

Parameters:

  • name (str)

  • available (Callable[[], bool])

  • target_kind (str)

  • nuts (Callable[..., NutsResult])

name: str
available: Callable[[], bool]
target_kind: str
nuts: Callable[..., NutsResult]
register_inference_backend(backend)[source]

Register (or replace) an inference backend under backend.name.

Parameters:

backend (InferenceBackend)

Return type:

None