blackjax.adaptation.staged_adaptation

blackjax.adaptation.staged_adaptation#

Staged warmup adaptation engine for HMC-family algorithms.

This module provides the staged_adaptation() engine and the build_schedule() function (previously in window_adaptation.py; re-exported from there for backward compatibility).

staged_adaptation() adapts step size and inverse mass matrix via the Stan warmup schedule for any algorithm whose kernel has signature:

kernel(rng_key, state, logdensity_fn, step_size, inverse_mass_matrix, **extra)

Supported: blackjax.nuts, blackjax.hmc, blackjax.mhmc, blackjax.barker, and others accepting the above contract. Excluded: RMHMC (kernel takes mass_matrix: Metric, not inverse_mass_matrix); GHMC/MEADS (kernel lacks inverse_mass_matrix); MCLMC (has own warmup); dynamic_hmc (init requires random_generator_arg). WindowAdaptationState in window_adaptation is an alias for StagedAdaptationState (same class object).

Notes

build_schedule is defined here (canonical location) and re-exported from window_adaptation for backward compatibility. Import from either module.

Classes#

StagedAdaptationState

Scan-carry state for the staged adaptation engine.

Functions#

build_schedule(→ list[tuple[int, bool]])

Return the schedule for Stan's warmup.

staged_adaptation(→ blackjax.base.AdaptationAlgorithm)

Adapt the step size and inverse mass matrix for HMC-family algorithms.

Module Contents#

class StagedAdaptationState[source]#

Scan-carry state for the staged adaptation engine.

Field names intentionally mirror the previous WindowAdaptationState fields so that any downstream code accessing adaptation info by field name (.ss_state, .imm_state, …) continues to work without modification.

WindowAdaptationState in blackjax.adaptation.window_adaptation is an alias of this type (WindowAdaptationState = StagedAdaptationState); both names refer to the same NamedTuple class object, so isinstance checks using either name are equivalent.

Parameters:

  • ss_state – Current state of the dual-averaging step-size adaptation.

  • imm_state – Current mass-matrix adaptation core state. One of MassMatrixAdaptationState or FisherMassMatrixAdaptationState. Typed Any here to avoid a hard dependency on the concrete types; the MetricCore protocol guarantees the right type at construction time.

  • step_size – Current (exponential-space) step-size estimate; read by the MCMC kernel at every scan step.

  • inverse_mass_matrix – Current inverse mass matrix; updated at each slow-window boundary and read by the MCMC kernel at every scan step.

ss_state: blackjax.adaptation.step_size.DualAveragingAdaptationState[source]#
imm_state: Any[source]#
step_size: float[source]#
inverse_mass_matrix: blackjax.types.Array[source]#
build_schedule(num_steps: int, initial_buffer_size: int = 75, final_buffer_size: int = 50, first_window_size: int = 25) → list[tuple[int, bool]][source]#

Return the schedule for Stan’s warmup.

The schedule below is intended to be as close as possible to Stan’s [stab]. The warmup period is split into three stages:

1. An initial fast interval to reach the typical set. Only the step size is adapted in this window. 2. “Slow” parameters that require global information (typically covariance) are estimated in a series of expanding intervals with no memory; the step size is re-initialized at the end of each window. Each window is twice the size of the preceding window. 3. A final fast interval during which the step size is adapted using the computed mass matrix.

Schematically:

` +---------+---+------+------------+------------------------+------+ |  fast   | s | slow |   slow     |        slow            | fast | +---------+---+------+------------+------------------------+------+ `

The distinction slow/fast comes from the speed at which the algorithms converge to a stable value; in the common case, estimation of covariance requires more steps than dual averaging to give an accurate value. See [stab] for a more detailed explanation.

Fast intervals are given the label 0 and slow intervals the label 1.

Parameters:

  • num_steps (int) – The number of warmup steps to perform.

  • initial_buffer_size (int) – The width of the initial fast adaptation interval.

  • first_window_size (int) – The width of the first slow adaptation interval.

  • final_buffer_size (int) – The width of the final fast adaptation interval.

Return type:

A list of tuples (window_label, is_middle_window_end).

staged_adaptation(algorithm, logdensity_fn: Callable, metric: str | blackjax.adaptation.metric_recipes.MetricRecipe | blackjax.adaptation.metric_recipes.MetricCore = 'welford_diag', *, max_grad_budget: int | None = None, n_chains: int = 1, imm_shrinkage_to_previous: float = 0.0, initial_inverse_mass_matrix: blackjax.types.Array | None = None, initial_step_size: float = 1.0, target_acceptance_rate: float = 0.8, adaptation_info_fn: Callable = return_all_adapt_info, integrator=mcmc.integrators.velocity_verlet, schedule_fn: Callable | None = None, initial_metric_state: Any = None, metric_telemetry: bool = False, telemetry_full_matrices: bool = False, **extra_parameters) → blackjax.base.AdaptationAlgorithm[source]#

Adapt the step size and inverse mass matrix for HMC-family algorithms.

The staged_adaptation() engine implements the same Stan warmup schedule as window_adaptation() but exposes a composable MetricCore interface for the mass-matrix adaptation component. The step-size dual-averaging and the stage schedule live in the HOST (this function); the mass-matrix estimation is fully delegated to the metric argument.

Parameters:

  • algorithm – A sampling algorithm whose kernel signature is (rng_key, state, logdensity_fn, step_size, inverse_mass_matrix, **extra_parameters), e.g. blackjax.nuts, blackjax.hmc, blackjax.mhmc. The algorithm’s build_kernel method is inspected to decide whether to pass an integrator.

  • logdensity_fn – The log density probability density function to sample.

  • metric –

    The mass-matrix adaptation specification. Accepts:

    • "auto" — the meta-adaptation controller (meta_adaptation). Automatically selects the diagonal vs low-rank path and the growing-window schedule. Requires max_grad_budget to be set. The emitted metric is always a LowRankInverseMassMatrix (with U=0, lam=1 when the controller stays diagonal — bit-equivalent to the diagonal metric).

      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. Use for exploration and algorithm development, not for production efficiency claims. A multi-chain escalation trigger (planned for v2) is expected to make the decision robust.

    • str — a registry name ("welford_diag" (default), "welford_dense", "fisher_diag"); looked up via lookup_recipe() and built with imm_shrinkage_to_previous and initial_inverse_mass_matrix.

    • MetricRecipe — built with imm_shrinkage_to_previous and initial_inverse_mass_matrix.

    • MetricCore — used directly as-is; imm_shrinkage_to_previous and initial_inverse_mass_matrix are ignored (closed over in the core).

  • max_grad_budget – Maximum total gradient budget (leapfrog evaluations). Required when metric="auto"; ignored otherwise. The meta-adaptation controller converts this to a warmup step count via a conservative divisor (see meta_adaptation). Passed as-is; use extract_meta_verdict() after warmup.run() to get the structured routing verdict and true gradient counts.

  • imm_shrinkage_to_previous – Pseudo-count controlling shrinkage of the per-window IMM toward the previous window’s IMM (Bayesian persistence). Default 0.0 reproduces Stan’s per-window-reset behavior exactly. Ignored when metric is a MetricCore.

  • initial_inverse_mass_matrix – Optional seed array for the initial inverse mass matrix. Ignored when metric is a MetricCore.

  • initial_step_size – Step size used to seed the dual-averaging adaptation.

  • target_acceptance_rate – Target Metropolis acceptance rate for step-size adaptation. Default 0.80 (Stan default).

  • metric_telemetry – Opt-in read-only observation of the metric-publication decision made at each slow-window boundary: the support actually consumed, the candidate metrics even when they are withheld, the raw-truth and escalation-applicability gate masks, and the step-size chronology across the boundary. Supported on both the single- and multi-chain metric="auto" paths; raises for a core that carries no publication record. Default False, which is a Python-time constant — the off path traces and computes exactly as before. Read the records with publication_adapt_info_fn() as adaptation_info_fn.

  • telemetry_full_matrices – Also carry the full candidate/deployed low-rank factors in each record. O(d*k) per record per step; off by default. Requires metric_telemetry=True.

  • adaptation_info_fn – Function to select the adaptation info returned at each step. See return_all_adapt_info() and get_filter_adapt_info_fn(). By default all information is saved — this can result in excessive memory usage if the information is unused.

  • integrator – The symplectic integrator passed to algorithm.build_kernel; only used if build_kernel accepts arguments. Defaults to velocity_verlet().

  • schedule_fn – Callable (num_steps: int) -> Array that returns a (num_steps, 2) array of (stage, is_window_end) pairs, or None (default) to use the path-appropriate default. When None and metric="auto", the default is build_growing_window_schedule() (nutpie’s proportional-to-tune, 1.5×-growing-window schedule). When None and any other metric, the default is build_schedule() (Stan’s fixed-absolute, 2×-doubling schedule). An explicit callable is always honored regardless of metric.

  • initial_metric_state – Optional pre-built mass-matrix adaptation core state. When not None, overrides the metric_core.init(n_dims) call at warmup start — the provided state is used as-is. The object must be a valid state for the chosen metric core (its inverse_mass_matrix field is unpacked into StagedAdaptationState immediately). Intended for callers that seed the initial state from external data (e.g., gradient-based diagonal-scale initialisation); None (the default) reproduces the standard identity/zero initialisation.

  • **extra_parameters – Algorithm-specific parameters forwarded to the MCMC kernel at every step, e.g. num_integration_steps for HMC/MHMC (divides budget when metric='auto') or num_max_steps for dynamic HMC.

Returns:

An AdaptationAlgorithm wrapping a run function with signature (rng_key, position, num_steps=1000) that returns (AdaptationResults, info).

Return type:

AdaptationAlgorithm

Notes

Wrap warmup.run(...) in blackjax.progress_bar() to display a progress bar, e.g. with blackjax.progress_bar(): warmup.run(...).

See also

blackjax.adaptation.window_adaptation.window_adaptation

Thin compatibility shim over this engine; preserves the old parameter interface exactly.

blackjax.adaptation.metric_recipes.REGISTRY

Registry of named MetricRecipe objects for the metric string argument.