blackjax.mcmc.laplace_hmc

blackjax.mcmc.laplace_hmc#

HMC on the Laplace-approximated marginal log-density with warm-starting.

Wraps the composable laplace_marginal_factory() in a standard BlackJAX three-layer sampler that carries the MAP latent variables theta_star through the chain. At each step theta_star is used as the warm-start hint for the L-BFGS solver at every leapfrog evaluation, so the optimizer needs only a handful of iterations when phi moves by a small amount.

The proposal strategy is swappable via build_proposal, giving two usable variants:

Alias	Proposal	Notes
`blackjax.laplace_hmc` `blackjax.laplace_mhmc`	endpoint + M-H full trajectory	default, standard HMC better ESS per gradient

Typical usage:

sampler = blackjax.laplace_hmc(
    log_joint, theta_init=jnp.zeros(n),
    step_size=0.1, inverse_mass_matrix=jnp.ones(d),
    num_integration_steps=10,
)
state = sampler.init(phi_init)
new_state, info = jax.jit(sampler.step)(rng_key, state)
# new_state.theta_star: MAP of theta at the accepted phi

# Multinomial variant (no rejection step, samples from full trajectory):
sampler = blackjax.laplace_mhmc(
    log_joint, theta_init=jnp.zeros(n),
    step_size=0.1, inverse_mass_matrix=jnp.ones(d),
    num_integration_steps=10,
)

Classes#

LaplaceHMCState

State of the Laplace-HMC sampler.

Functions#

`init`(→ LaplaceHMCState)	Create an initial `LaplaceHMCState`.
`build_kernel`(→ Callable)	Build the Laplace-HMC kernel.
`as_top_level_api`(→ blackjax.base.SamplingAlgorithm)	HMC on the Laplace-approximated marginal log-density.

Module Contents#

class LaplaceHMCState[source]#

State of the Laplace-HMC sampler.

position: Current hyperparameter position phi. Can be any PyTree.
logdensity: Current value of the Laplace log-marginal log p̂(phi | y).
logdensity_grad: Gradient of log p̂(phi | y) w.r.t. phi. Same PyTree structure as position.
theta_star: MAP of the latent variables at the current phi, i.e. theta*(phi) = argmax_theta log_joint(theta, phi). Carried through the chain so the next L-BFGS solve can warm-start from here.

position: blackjax.types.ArrayTree[source]#

logdensity: float[source]#

logdensity_grad: blackjax.types.ArrayTree[source]#

theta_star: blackjax.types.ArrayTree[source]#

init(position: blackjax.types.ArrayLikeTree, laplace: blackjax.mcmc.laplace_marginal.LaplaceMarginal) → LaplaceHMCState[source]#

Create an initial LaplaceHMCState.

Runs L-BFGS from cold start to find theta*(position), then evaluates the Laplace log-marginal and its gradient.

Parameters:

position – Initial hyperparameter value phi.
laplace – A LaplaceMarginal instance returned by laplace_marginal_factory().

build_kernel(integrator: Callable = integrators.velocity_verlet, divergence_threshold: float = 1000, build_proposal: Callable = hmc.hmc_proposal) → Callable[source]#

Build the Laplace-HMC kernel.

Parameters:

integrator – Symplectic integrator used for the HMC trajectory.
divergence_threshold – Energy difference above which a transition is declared divergent.
build_proposal – Proposal builder. Defaults to hmc_proposal() (endpoint + M-H). Pass multinomial_hmc_proposal() for multinomial trajectory sampling (blackjax.laplace_mhmc).

Return type:

A kernel (rng_key, state, laplace, step_size, inverse_mass_matrix, num_integration_steps) -> (LaplaceHMCState, LaplaceHMCInfo).

as_top_level_api(log_joint_fn: Callable, theta_init: blackjax.types.ArrayLikeTree, step_size: float, inverse_mass_matrix: blackjax.mcmc.metrics.MetricTypes, num_integration_steps: int, *, divergence_threshold: int = 1000, integrator: Callable = integrators.velocity_verlet, build_proposal: Callable = hmc.hmc_proposal, **optimizer_kwargs) → blackjax.base.SamplingAlgorithm[source]#

HMC on the Laplace-approximated marginal log-density.

For a hierarchical model log p(theta, phi, y), integrates out the latent variables theta via the Laplace approximation and runs HMC on the resulting marginal over the hyperparameters phi.

Gradients w.r.t. phi are computed via the implicit function theorem (jax.lax.custom_root()) — the L-BFGS iterations are not unrolled. theta*(phi) is warm-started from the previous MCMC state, reducing the number of L-BFGS iterations needed at each leapfrog step.

Parameters:

log_joint_fn – (theta, phi) -> float. The full log joint log p(theta, phi, y). Both arguments may be arbitrary PyTrees. Must be at least C³ in theta.
theta_init – Initial guess for theta. Fixes the PyTree structure for all calls.
step_size – HMC leapfrog step size.
inverse_mass_matrix – Inverse mass matrix for HMC (1-D array for diagonal, scalar for isotropic).
num_integration_steps – Number of leapfrog steps per HMC transition.
divergence_threshold – Absolute energy difference above which a transition is declared divergent. Default 1000.
integrator – Symplectic integrator. Default: velocity Verlet.
build_proposal – Proposal builder. Defaults to hmc_proposal() (endpoint + M-H). Pass multinomial_hmc_proposal() for multinomial trajectory sampling; this is what blackjax.laplace_mhmc uses.
**optimizer_kwargs – Forwarded to minimize_lbfgs(). Useful keys: maxiter (default 30), gtol, ftol.

Returns:

A SamplingAlgorithm whose step returns a
LaplaceHMCState (with theta_star field) and
LaplaceHMCInfo (includes all
standard HMCInfo fields plus L-BFGS
diagnostics lbfgs_iter_num, lbfgs_error, lbfgs_converged,
lbfgs_hit_maxiter).

Examples

sampler = blackjax.laplace_hmc(
    log_joint, theta_init=jnp.zeros(n_latent),
    step_size=0.1, inverse_mass_matrix=jnp.ones(d_phi),
    num_integration_steps=10, maxiter=100,
)
state = sampler.init(phi_init)
new_state, info = jax.jit(sampler.step)(rng_key, state)
print(new_state.theta_star)   # MAP latent at the new phi