blackjax.adaptation.mclmc_adaptation

blackjax.adaptation.mclmc_adaptation#

Algorithms to adapt the MCLMC kernel parameters, namely step size and L.

Classes#

MCLMCAdaptationState

Represents the tunable parameters for MCLMC adaptation.

Functions#

`mclmc_find_L_and_step_size`(mclmc_kernel, num_steps, ...)	Finds the optimal value of the parameters for the MCLMC algorithm.
`make_L_step_size_adaptation`(kernel, logdensity_fn, ...)	Adapts the stepsize and L of the MCLMC kernel. Designed for unadjusted MCLMC
`make_adaptation_L`(kernel, logdensity_fn, frac, l_factor)	determine L by the autocorrelations (around 10 effective samples are needed for this to be accurate)
`handle_nans`(previous_state, next_state, step_size, ...)	Adaptation-level NaN handler.

Module Contents#

class MCLMCAdaptationState[source]#

Represents the tunable parameters for MCLMC adaptation.

L: The momentum decoherent rate for the MCLMC algorithm.
step_size: The step size used for the MCLMC algorithm.
inverse_mass_matrix: A matrix used for preconditioning.

L: float[source]#

step_size: float[source]#

inverse_mass_matrix: float[source]#

mclmc_find_L_and_step_size(mclmc_kernel, num_steps, state, rng_key, logdensity_fn=None, frac_tune1=0.1, frac_tune2=0.1, frac_tune3=0.1, desired_energy_var=0.0005, trust_in_estimate=1.5, num_effective_samples=150, diagonal_preconditioning=True, params=None, l_factor=0.4)[source]#

Finds the optimal value of the parameters for the MCLMC algorithm.

Parameters:

mclmc_kernel – The kernel function built by mclmc.build_kernel. Its call signature must be kernel(rng_key, state, logdensity_fn, inverse_mass_matrix, L, step_size), matching the standard BlackJAX kernel pattern.
num_steps – The number of MCMC steps that will subsequently be run, after tuning.
state – The initial state of the MCMC algorithm.
rng_key – The random number generator key.
logdensity_fn – The log-density function of the target distribution.
frac_tune1 – The fraction of tuning for the first step of the adaptation.
frac_tune2 – The fraction of tuning for the second step of the adaptation.
frac_tune3 – The fraction of tuning for the third step of the adaptation.
desired_energy_var – The desired energy variance for the MCMC algorithm.
trust_in_estimate – The trust in the estimate of optimal stepsize.
num_effective_samples – The number of effective samples for the MCMC algorithm.
diagonal_preconditioning – Whether to do diagonal preconditioning (i.e. a mass matrix)
params – Initial params to start tuning from (optional)
l_factor – The factor scaling the estimated autocorrelation length to obtain momentum decoherence length L.

Returns:

final_state – The final integrator state after the three tuning phases.
final_params – An MCLMCAdaptationState containing the adapted L, step_size, and inverse_mass_matrix.
total_num_tuning_integrator_steps – The total number of integrator steps consumed across all three tuning phases (frac_tune1 + frac_tune2 + frac_tune3 of num_steps).

Example

kernel = blackjax.mcmc.mclmc.build_kernel(integrator=integrator)

(
    blackjax_state_after_tuning,
    blackjax_mclmc_sampler_params,
    num_tuning_steps,
) = blackjax.mclmc_find_L_and_step_size(
    mclmc_kernel=kernel,
    logdensity_fn=logdensity_fn,
    num_steps=num_steps,
    state=initial_state,
    rng_key=tune_key,
    diagonal_preconditioning=preconditioning,
)

Notes

Live divergence monitoring (jax-tap >= 0.3.0)

The internal tuning scan exposes a per-step divergence flag as its ys output (True = divergence on that step). Users who install jax-tap >= 0.3.0 can observe this stream with no changes to BlackJAX:

import jaxtap  # pip install "jax-tap>=0.3.0"

with jaxtap.record(
    select_ys=lambda ys: ys[0],  # the single divergence-flag leaf
    alert_ys=lambda e: "divergence" if e.value else None,
    alert_ys_once=True,  # one stderr line then silence; drop for per-step
) as rec:
    state, params, _ = blackjax.mclmc_find_L_and_step_size(
        mclmc_kernel=kernel, num_steps=N, state=init_state,
        rng_key=key, logdensity_fn=logdensity_fn,
    )
divergence_steps = [
    e.step for e in rec.events if e.kind == "output" and e.value
]

Checking for degenerate warmup

BlackJAX does not emit runtime warnings; checking is the user’s responsibility.

Before calling — verify the initial gradient is finite:

from jax.flatten_util import ravel_pytree
ok = jnp.all(jnp.isfinite(ravel_pytree(state.logdensity_grad)[0]))
# finite logdensity + non-finite gradient = model/solver/support issue (#973)

After calling — a collapsed warmup leaves step_size orders of magnitude below the posterior scale; healthy and frozen runs differ by ~6 orders:

ratio = final_params.step_size * num_steps / final_params.L
# ratio ≈ 1 → healthy;  ratio << 1 → likely frozen

make_L_step_size_adaptation(kernel, logdensity_fn, dim, frac_tune1, frac_tune2, diagonal_preconditioning, desired_energy_var=0.001, trust_in_estimate=1.5, num_effective_samples=150)[source]#: Adapts the stepsize and L of the MCLMC kernel. Designed for unadjusted MCLMC

make_adaptation_L(kernel, logdensity_fn, frac, l_factor)[source]#: determine L by the autocorrelations (around 10 effective samples are needed for this to be accurate)

handle_nans(previous_state, next_state, step_size, step_size_max, kinetic_change, kernel_nonans, key)[source]#

Adaptation-level NaN handler.

If the kernel reported a divergence (via its truthful info.nonans after #969 fix), reduce step_size_max and return the pre-step state. The kernel’s own handle_nans already sanitises next_state for both divergence signatures:

Case-1: NaN position or momentum (position overshoot through a hard boundary).
Case-2: finite position + momentum but NaN logdensity (dominant under velocity_verlet at moderate overshoot on bounded targets).

Parameters:: kernel_nonans – info.nonans from the MCLMC kernel — truthful after the #969 fix.
Returns:: True when the step was clean (no divergence and finite energy change).
Return type:: success