blackjax.adaptation.metric_estimators

blackjax.adaptation.metric_estimators#

Pure metric estimator functions for the low-rank metric adaptation layer.

Each function is a pure transformation: explicit arrays in → metric representation out. No buffer state, no scheduling logic, no side effects. All functions are JAX-traceable and safe to use inside jax.lax.scan / jax.vmap provided that max_rank (where applicable) is a static Python integer (required to determine output shapes).

Source lineage (estimator extraction history; the private _compute_low_rank_metric helper shared by the fisher path now resides in this module, relocated from low_rank_adaptation):

Estimator

Extracted from

fisher_score_low_rank()

low_rank_adaptation._compute_low_rank _metric + _spd_mean

draws_singular_value_low_rank()

mclmc_lrd_adaptation ._extract_lrd_from_samples

sample_covariance_eigh_low`() ``_rank`

meads_adaptation ._lrd_from_accumulated_covariance

welford_diagonal() welford_dense()

mass_matrix.welford_algorithm (thin scan-wrappers; algorithm untouched)

fisher_score_diagonal()

branch b197f1e2 mass_matrix._fisher_diagonal_inverse _mass (dormant; wiring stays dormant)

sample_variance_diagonal()

mclmc_adaptation.py:342 and adjusted_mclmc_adaptation.py:375 (verbatim duplicate; one function here)

Composition note: estimators, data-feeding policy, and schedule are co-adapted package components. The functions here are the estimator component only — callers supply the buffer view. Gating logic (support gates, fraction-window checks, 2·d thresholds) belongs in the caller; it is explicitly not implemented here (docstrings note the relevant gates for each estimator).

Registration: these functions are module-public (importable from blackjax.adaptation.metric_estimators) but are NOT exported at the blackjax top-level. Top-level export and consumer re-wiring are deliberate follow-up work.

Functions#

eigenvalue_informativeness(→ blackjax.types.Array)

Informativeness of each eigenvalue: how far it deviates from isotropic.

select_top_eigenvalues_by_informativeness(...)

Select the top-max_rank eigenpairs by eigenvalue_informativeness().

fisher_score_low_rank(...)

Fisher-divergence-minimising low-rank inverse mass matrix.

draws_singular_value_low_rank(...)

Draws-only SVD low-rank inverse mass matrix (MCLMC Scheme A estimator).

sample_covariance_eigh_low_rank(...)

Chan/Welford-accumulated covariance → low-rank metric via eigh.

welford_diagonal(→ blackjax.types.Array)

Diagonal (per-coordinate) sample variance via Welford's algorithm.

welford_dense(→ blackjax.types.Array)

Dense sample covariance via Welford's algorithm.

fisher_score_diagonal_from_moments(→ blackjax.types.Array)

Core math of the Fisher-diagonal estimator from pre-computed variances.

fisher_score_diagonal(→ blackjax.types.Array)

Fisher-divergence-minimising diagonal inverse mass matrix.

sample_variance_diagonal(→ blackjax.types.Array)

Coordinate-wise population variance from raw draws.

Module Contents#

eigenvalue_informativeness(eigenvalues: blackjax.types.Array) → blackjax.types.Array[source]#

Informativeness of each eigenvalue: how far it deviates from isotropic.

Computes \(|\lambda - 1|\) for each entry of eigenvalues. Directions with large informativeness deviate the most from the identity metric and benefit the most from low-rank preconditioning.

This idiom is used identically in:

Parameters:

eigenvalues – Shape (q,). Eigenvalues of a correlation (or preconditioned) matrix. An isotropic Gaussian has all eigenvalues equal to 1 and informativeness 0.

Returns:

\(|\lambda_i - 1|\) for each \(\lambda_i\).

Return type:

Array, shape (q,)

select_top_eigenvalues_by_informativeness(eigenvalues: blackjax.types.Array, eigenvectors: blackjax.types.Array, max_rank: int, *, tail_handling: Literal['mask_pad', 'raw'] = 'mask_pad', cutoff: float = 2.0) → tuple[blackjax.types.Array, blackjax.types.Array][source]#

Select the top-max_rank eigenpairs by eigenvalue_informativeness().

Two production consumers use this pattern with divergent tail handling:

tail_handling="mask_pad" (default)

Used by fisher_score_low_rank(). Eigenvalues inside the uninformative band \([1/\text{cutoff},\ \text{cutoff}]\) are masked to 1 (they carry no preconditioning benefit), making the corresponding direction a no-op in the metric. When the input subspace has fewer than max_rank eigenvectors (q < max_rank, which arises when d < 2 · max_rank), the output is zero-padded to the requested shape — zero columns in U and lam=1 entries — so the return shape is always (d, max_rank) / (max_rank,). Matches _compute_low_rank_metric (this module) steps 8–9.

tail_handling="raw"

Used by draws_singular_value_low_rank() and sample_covariance_eigh_low_rank(). The top-k pairs are returned as-is: no informativeness masking, no padding. The caller must ensure max_rank <= q; behaviour is undefined (JAX will silently return q columns) when max_rank > q. Matches mclmc_lrd_adaptation._extract_lrd_from_samples :271–275 and meads_adaptation._lrd_from_accumulated_covariance :309–312.

Tie-break divergence: the two modes differ not only in masking but also in sort stability for exact |λ−1| ties: "mask_pad" uses argsort(-scores) (ascending original-index among ties; Fisher consumer convention); "raw" uses argsort(scores)[::-1] (descending original-index; both raw-source conventions). On real eigh/svd spectra bit-exact ties do not occur — even truly degenerate inputs return ulp-broken eigenvalues — so the difference is behaviorally inert on continuous data; the raw path preserves byte-fidelity to its sources nonetheless.

Parameters:

  • eigenvalues – Shape (q,). Must be the eigenvalues associated with the columns of eigenvectors.

  • eigenvectors – Shape (d, q). Columns are the eigenvectors.

  • max_rank – Number of eigenpairs to return. Must be a static Python integer (determines the output shape).

  • tail_handling – See above. Default "mask_pad" matches the Fisher-score consumer.

  • cutoff – Only used when tail_handling="mask_pad". Eigenvalues in \([1/\text{cutoff},\ \text{cutoff}]\) are masked to 1. Default 2.0 matches nutpie’s c=2.

Returns:

  • U_out (Array, shape (d, max_rank) (mask_pad) or (d, ≤max_rank) (raw)) – Selected (and possibly masked/padded) eigenvectors.

  • lam_out (Array, shape (max_rank,) or (≤max_rank,)) – Selected (and possibly masked) eigenvalues.

fisher_score_low_rank(draws: blackjax.types.Array, grads: blackjax.types.Array, max_rank: int, *, gamma: float = 1e-05, cutoff: float = 2.0) → blackjax.mcmc.metrics.LowRankInverseMassMatrix[source]#

Fisher-divergence-minimising low-rank inverse mass matrix.

Implements Steps 1–9 of Algorithm 1 of [SCC26], following the nutpie reference implementation (nuts-rs src/transform/adapt/low_rank.rs estimate_mass_matrix).

The inverse mass matrix has the form

\[M^{-1} = \operatorname{diag}(\sigma) \bigl(I + U(\Lambda - I)U^\top\bigr) \operatorname{diag}(\sigma)\]

where \(\sigma\), \(U\), \(\Lambda\) minimise the sample Fisher divergence from \(\{(x_i, \nabla \log p(x_i))\}\).

Extracted from: blackjax.adaptation.low_rank_adaptation ._compute_low_rank_metric + ._spd_mean (main @ 532631c1).

Key asymmetry vs draws_singular_value_low_rank(): this estimator uses BOTH draws and score gradients and applies γ-regularisation plus cutoff-masking on the eigenvalues. draws_singular_value_low_rank() uses draws only and applies neither — the docstring notes this divergence explicitly.

Diagonal scale σ = (Var[x] / Var[∇ log p])^{1/4} is the per-coordinate optimal scale (paper §3.1); clipped to [1e-20, 1e20] (nutpie range).

AIRM geometric mean Σ = C_x # C_a^{-1} (AIRM = affine-invariant Riemannian metric): Theorem 2.3 / Eq. 9 of [SCC26] — the regularised optimal inverse mass matrix is the geometric mean of the draw covariance with the inverse score covariance.

γ-regularisation C = P P^T / γ + I (nutpie convention: the unnormalised sum-of-outer-products divided by γ directly, no n scaling). Influence fades as n grows (Theorem 2.4).

Cutoff masking: eigenvalues in [1/cutoff, cutoff] are set to 1 (no preconditioning benefit); default cutoff=2 matches nutpie.

dtype promotion (round-9 schedule-port audit): promotes internally to float64 when jax_enable_x64 is active, regardless of the chain’s working dtype. Returns in the caller’s original dtype.

G-layer note: the 2·d support gate (n ≥ 2·d before the LR estimate is trusted) and any fraction-window guard are G-layer concerns and are NOT implemented here.

Parameters:

  • draws – Shape (n, d). Chain positions (all rows must be valid samples — no zero-padding).

  • grads – Shape (n, d). Log-density gradients at the corresponding draws.

  • max_rank – Maximum number of eigenvectors in the low-rank correction. Must be a static Python integer (determines output shape).

  • gamma – Regularisation scale (nutpie convention). Default 1e-5 matches nutpie’s LowRankSettings::default.

  • cutoff – Eigenvalue cutoff for informativeness masking. Default 2.0 matches nutpie’s c=2.

Returns:

(sigma, U, lam) with shapes (d,), (d, max_rank), (max_rank,).

Return type:

LowRankInverseMassMatrix

Notes

The optimal translation μ* = x̄ + σ² ⊙ ᾱ (paper §3.2) is an adaptation-layer output, not part of the metric. Compute it separately from the per-draw means if needed by the warmup wiring (a warmup-wiring concern, deliberately outside this estimator).

draws_singular_value_low_rank(draws: blackjax.types.Array, max_rank: int) → blackjax.mcmc.metrics.LowRankInverseMassMatrix[source]#

Draws-only SVD low-rank inverse mass matrix (MCLMC Scheme A estimator).

Estimates the low-rank inverse mass matrix from the SVD of centred, standardised draws. No gradient information is required.

Extracted from: blackjax.adaptation.mclmc_lrd_adaptation ._extract_lrd_from_samples (main @ 532631c1).

Key asymmetry vs fisher_score_low_rank():

  • Draws only — no score gradients.

  • No regularisation (no γ): the covariance is estimated from raw outer products without a P P^T / γ + I regularisation term.

  • No masking: eigenvalues are returned as-is (tail_handling="raw" in select_top_eigenvalues_by_informativeness()); non-informative eigenvalues are NOT masked to 1. This is a deliberate design choice in the MCLMC-LRD pilot estimator: the raw eigenspectrum is what drives the effective condition number diagnostics downstream.

This asymmetry is preserved faithfully here; see the module docstring’s source-lineage table for context.

G-layer note: the Geyer-ESS support gate (k_used = min(k, ⌊n_eff/2⌋), mclmc_lrd_adaptation.py:636) and the n ≥ 2·d threshold are G-layer concerns. Ensure max_rank ≤ min(n, d) before calling (behaviour is undefined otherwise; JAX will return fewer than max_rank columns when max_rank > min(n, d)).

Parameters:

  • draws – Shape (n, d). All rows must be valid samples (no zero-padding).

  • max_rank – Number of eigenpairs to return. Must be a static Python integer and satisfy max_rank ≤ min(n, d).

Returns:

(sigma, U, lam) where sigma is the per-coordinate standard deviation and lam are the raw SVD eigenvalues of the sample correlation matrix (not masked to 1 for near-unity values).

Return type:

LowRankInverseMassMatrix

Notes

The full eigenspectrum needed for _kappa_eff_pilot() diagnostics is NOT returned here — that function operates on the full sorted spectrum (lam_all_sorted). If you need the full spectrum, call _extract_lrd_from_samples directly (G-layer concern, not part of the pure estimator).

sample_covariance_eigh_low_rank(m2: blackjax.types.Array, count: blackjax.types.Array | int, max_rank: int) → blackjax.mcmc.metrics.LowRankInverseMassMatrix[source]#

Chan/Welford-accumulated covariance → low-rank metric via eigh.

Extracts a low-rank inverse mass matrix from a Chan-parallel-accumulated sum of squared deviations matrix (MEADS / MCLMC-LRD Scheme-B estimator).

Extracted from: blackjax.adaptation.meads_adaptation ._lrd_from_accumulated_covariance (main @ 532631c1, landed via #954).

The input m2 is the accumulated Chan-parallel M2 matrix (shape (d, d)):

\[M_2 = \sum_{i=1}^n (x_i - \bar{x}_n)(x_i - \bar{x}_n)^T\]

which gives the Bessel-corrected sample covariance as C = M_2 / (n - 1). The correlation matrix is then R = D^{-1/2} C D^{-1/2} where D = diag(C), and eigh of R gives the eigenbasis.

Tail handling: raw (no cutoff masking, no zero-padding) — matching the MEADS implementation. The G-layer (meads_adaptation’s low_rank_window_fraction / 2·d gate) ensures the estimate is only used when the accumulated support suffices.

G-layer note: the fraction-window gate (only accumulate draws in the second half of the adaptation window, low_rank_window_fraction=0.5) and the 2·d support threshold are G-layer concerns; they gate when to call this estimator, not what it computes. Callers must ensure count reflects enough effective support before calling.

Parameters:

  • m2 – Shape (d, d). Accumulated Chan-Welford sum of squared deviations (NOT divided by count; this function applies the Bessel correction / max(count - 1, 1) internally).

  • count – Total number of samples accumulated into m2. May be a traced JAX integer (safe inside jax.lax.scan).

  • max_rank – Number of eigenpairs to return. Must be a static Python integer and satisfy max_rank ≤ d.

Returns:

(sigma, U, lam) with shapes (d,), (d, max_rank), (max_rank,).

Return type:

LowRankInverseMassMatrix

welford_diagonal(draws: blackjax.types.Array) → blackjax.types.Array[source]#

Diagonal (per-coordinate) sample variance via Welford’s algorithm.

Thin scan-wrapper over blackjax.adaptation.mass_matrix.welford_algorithm() with is_diagonal_matrix=True. Computes the Bessel-corrected sample variance \(s^2_i = \frac{1}{n-1}\sum_{j=1}^n (x_{ji} - \bar{x}_i)^2\) for each coordinate i.

Why a wrapper instead of direct ``jnp.var``: having a single estimator import surface ensures future callers can swap to streaming Welford (e.g., for online adaptation) without changing call sites. The algorithm itself is NOT moved or duplicated — only the call site is unified here.

Source: blackjax.adaptation.mass_matrix.welford_algorithm (is_diagonal_matrix=True).

Parameters:

draws – Shape (n, d).

Returns:

Bessel-corrected per-coordinate sample variance (= diagonal of the sample covariance matrix).

Return type:

Array, shape (d,)

welford_dense(draws: blackjax.types.Array) → blackjax.types.Array[source]#

Dense sample covariance via Welford’s algorithm.

Thin scan-wrapper over blackjax.adaptation.mass_matrix.welford_algorithm() with is_diagonal_matrix=False. Computes the Bessel-corrected sample covariance matrix.

Source: blackjax.adaptation.mass_matrix.welford_algorithm (is_diagonal_matrix=False).

Parameters:

draws – Shape (n, d).

Returns:

Bessel-corrected sample covariance matrix.

Return type:

Array, shape (d, d)

fisher_score_diagonal_from_moments(variance: blackjax.types.Array, gradient_variance: blackjax.types.Array) → blackjax.types.Array[source]#

Core math of the Fisher-diagonal estimator from pre-computed variances.

Computes \(\sigma^2 = \sqrt{\mathrm{Var}[x] / \mathrm{Var}[\nabla \log p]}\) per coordinate — the same formula as fisher_score_diagonal() but operating on already-computed per-coordinate variances rather than raw draw arrays.

This entry point is intended for callers that accumulate moments online (e.g. via _FisherMomentBlock) and want to avoid materialising the full draw array. The caller is responsible for supplying Bessel-corrected (or otherwise normalised) variances; the ratio is invariant to a shared n vs n-1 factor so either convention is acceptable for the diagonal estimator.

Near-zero gradient protection (identical to fisher_score_low_rank() and fisher_score_diagonal()): gradient_variance is floored at 1e-10 before division, and the result is clipped to nutpie’s [1e-20, 1e20] range before squaring.

Pairing insensitivity: this function consumes only the marginal variances Var[x_i] and Var[∇log p_i] per coordinate. The estimator is therefore insensitive to which draw is paired with which gradient within a batch — any per-batch pairing permutation produces the same variance and gradient_variance inputs and hence the same output. If cross-moment information (e.g. draw-grad covariance) is required, a future extension to this signature must add those moments explicitly.

Planned extension note: this entry point is intentionally separate so that future updates to the estimator (e.g. adding draw-grad cross moments) can extend the from_moments signature without changing the raw-draws wrapper fisher_score_diagonal().

Parameters:

  • variance – Shape (d,). Per-coordinate position variance \(\mathrm{Var}[x]\). Must be non-negative; typically Bessel-corrected.

  • gradient_variance – Shape (d,). Per-coordinate log-density-gradient variance \(\mathrm{Var}[\nabla \log p]\). Must be non-negative; floored at 1e-10 internally.

Returns:

Diagonal inverse mass matrix \(\sigma^2 = \sigma_{\text{clip}}^2\) where \(\sigma_{\text{clip}} = \operatorname{clip}(\sigma,\, 10^{-20},\, 10^{20})\) and \(\sigma = (\mathrm{Var}[x] / \max(\mathrm{Var}[\nabla \log p],\, 10^{-10}))^{1/4}\).

The clip is on the scale \(\sigma\), not on \(\sigma^2\), so the returned values span \([10^{-40}, 10^{40}]\). Under float32, \(\sigma^2\) overflows to inf when \(\sigma \approx 10^{20}\) (float32 max \(\approx 3.4 \times 10^{38}\), so \(10^{40}\) overflows); the caller must handle this if float32 inputs can drive \(\sigma\) to its clip boundary. This matches fisher_score_low_rank()’s sigma clip and is not changed here to preserve numerical consistency between the two estimators.

Return type:

Array, shape (d,)

fisher_score_diagonal(draws: blackjax.types.Array, grads: blackjax.types.Array) → blackjax.types.Array[source]#

Fisher-divergence-minimising diagonal inverse mass matrix.

Computes \(\sigma^2 = \sqrt{\mathrm{Var}[x] / \mathrm{Var}[\nabla \log p]}\), the per-coordinate diagonal estimator of [SCC26].

This is the diagonal-only analogue of fisher_score_low_rank(): that function’s diagonal scale \(\sigma = (\mathrm{Var}[x] / \mathrm{Var}[\nabla \log p])^{1/4}\) gives the corresponding inverse mass matrix as \(\sigma^2 = \sqrt{\mathrm{Var}[x] / \mathrm{Var}[\nabla \log p]}\) when the low-rank correction (U, lam) is dropped.

Extracted from: branch b197f1e2 (feat/window-adaptation-fisher-diag, 2026-07-04), blackjax.adaptation.mass_matrix._fisher_diagonal_inverse_mass.

Near-zero gradient protection (same as fisher_score_low_rank()): Var[∇ log p] is floored at 1e-10 before division, and the result is clipped to nutpie’s [1e-20, 1e20] range before squaring.

Note on variance convention: Var[x] and Var[∇ log p] are computed with the same normalisation (Bessel-corrected, n-1 divisor, via welford_diagonal()); the ratio is invariant to a shared n vs n-1 factor (branch b197f1e2 commit message, verified).

Implementation: thin wrapper over fisher_score_diagonal_from_moments() — computes Bessel-corrected per-coordinate variances via two welford_diagonal() scans, then delegates all arithmetic to the from-moments entry point.

Parameters:

  • draws – Shape (n, d). Chain positions.

  • grads – Shape (n, d). Log-density gradients at the corresponding draws.

Returns:

Diagonal inverse mass matrix \(\sigma^2\).

Return type:

Array, shape (d,)

sample_variance_diagonal(draws: blackjax.types.Array) → blackjax.types.Array[source]#

Coordinate-wise population variance from raw draws.

Computes \(\hat{\sigma}^2_i = \mathbb{E}[x_i^2] - \mathbb{E}[x_i]^2\) (population variance, no Bessel correction) for each coordinate i.

Extracted from (verbatim duplicate):

  • blackjax.adaptation.mclmc_adaptation, L_step_size_adaptation :341–342: variances = x_squared_average - jnp.square(x_average) where x_average = E[x] and x_squared_average = E[x^2] are streaming step-size-weighted averages.

  • blackjax.adaptation.adjusted_mclmc_adaptation, adjusted_mclmc_find_L_and_step_size :374–375: identical formula.

Both inline occurrences use the result as inverse_mass_matrix = variances, i.e., the population variance IS the diagonal IMM. The streaming-average formulation (weighted by step size via incremental_value_update) is an adaptation-layer scheduling concern; the pure estimator here takes the batch of accumulated draws directly.

Population (not Bessel-corrected) to match the MCLMC streaming form: the weighted incremental average avg ← avg + w·(x − avg) / Σw converges to \(\mathbb{E}[x]\) under uniform weights, not the unbiased \(n/(n-1)\) form.

Parameters:

draws – Shape (n, d). The draws accumulated over the estimation window.

Returns:

Per-coordinate population variance \(\mathbb{E}[x^2] - \mathbb{E}[x]^2\).

Return type:

Array, shape (d,)