blackjax.adaptation.meta

blackjax.adaptation.meta#

Meta-adaptation controller for the HMC-family warmup.

At each window boundary the controller computes two signals: (1) held-out score-linearity R² — the curvature gate (funnel R²≈0.007 vs ≥0.54 for all metric-fixable classes); (2) S_gap(k) = λ₁/λ_{k+1} of the Welford-whitened residual — the magnitude predictor (Spearman 1.0 with measured rank-k payoff). Escalate diagonal → rank-k iff R² ≥ _R_MIN AND S_gap ≥ _S_MIN AND stable over two consecutive windows AND budget deadline clear. Growing-window schedule (nutpie-style) is the default; AIRM-velocity early exit is advisory (converged_at_step records where stopping would have helped).

Warning

metric="auto" is experimental (v1). The low-rank escalation is not robustly calibrated at high dimension: when the residual spectrum’s dominant structure sits near the detection boundary, whether the controller escalates can depend on the random seed used for sampling. Use metric="auto" for exploration and algorithm development, not for production efficiency claims. A multi-chain escalation trigger (planned for v2) is expected to make the escalation decision robust across seeds.

Dtype note: the composed estimator _compute_low_rank_metric produces numerically indefinite metrics under float32 (~98% of runs). Enable x64 via jax.config.update("jax_enable_x64", True) for production use.

See blackjax.adaptation.metric_recipes for the MetricCore protocol and blackjax.adaptation.staged_adaptation for the host engine.

Submodule layout#

_state

State NamedTuples: MetaAdaptationCoreState, MultiChainMetaAdaptationCoreState, MetaAdaptationVerdict.

_calibration

All module constants and swappable calibration functions.

_signals

Single-chain signal computation (spectrum, S_gap, R², mixing).

_detection

Multi-chain detection statistics (within-chain, between-chain, LOO, gap-stat).

_router

Router functions (geometric-mean scale, projected-gain R², PC-centered pool).

_schedule

Multi-chain growing-window schedule.

builders

Core builder functions: build_meta_adaptation_core, build_multi_chain_meta_core.

verdict

Post-run verdict extractors: extract_meta_verdict, extract_multi_chain_verdict. Opt-in publication telemetry (single-chain): MetricPublicationRecord, publication_adapt_info_fn, extract_publication_chronology, decode_gates.

Submodules#