Skip to content

MCMC Configurations

MCMCInference is configured via MCMC config dataclasses. These specify sampler family, backend source, and algorithm hyperparameters.

BaseMCMCConfig dataclass

Shared configuration options inherited by all MCMC configs.

You do not instantiate this class directly; use one of the concrete subclasses (NUTSConfig, HMCConfig, AdaptiveMetropolisConfig, SGLDConfig, MALAConfig, AdjustedMCLMCDynamicConfig).

Attributes:

Name Type Description
num_samples int

Number of post-warmup samples to return.

num_warmup int

Number of warmup/burn-in transitions.

num_chains int

Number of Markov chains to run in parallel.

mcmc_source MCMCSource

Backend library used for inference. Supported values are "numpyro" and "blackjax".

init_strategy callable

NumPyro initialization strategy used when constructing unconstrained initial parameters.

NUTSConfig dataclass

Bases: BaseMCMCConfig

No-U-Turn Sampler (NUTS) configuration.

Attributes:

Name Type Description
target_acceptance_rate float

Target acceptance probability used during warmup. Must lie strictly between zero and one.

HMCConfig dataclass

Bases: BaseMCMCConfig

Hamiltonian Monte Carlo (HMC) configuration.

Attributes:

Name Type Description
step_size float

Integrator step size used by the leapfrog solver. Ignored when adapt=True (step size is tuned during warmup).

num_steps int

Number of leapfrog steps per HMC proposal.

adapt bool

Whether to tune step size and mass matrix during warmup. Defaults to True. Set to False to use a fixed step_size and identity mass matrix (useful when warmup is expensive or the step size is known in advance).

AdaptiveMetropolisConfig dataclass

Bases: BaseMCMCConfig

Adaptive random-walk Metropolis-within-Gibbs configuration.

One transition updates each flattened unconstrained coordinate in order. Each coordinate update uses a one-dimensional Gaussian random-walk proposal that is accepted or rejected before the next coordinate is visited; the sampler does not make a single joint multivariate-normal proposal. Proposal scales adapt during warmup toward the requested acceptance rate and remain fixed while retained samples are generated.

This sampler is currently implemented by the BlackJAX integration only.

Attributes:

Name Type Description
initial_proposal_scale ArrayLike

Positive scalar proposal scale, or one positive scale per flattened unconstrained coordinate.

target_acceptance_rate float

Shared target acceptance rate applied to each coordinate update.

adaptation_rate float

Exponent in the diminishing adaptation step n_iter ** -adaptation_rate.

max_adaptation float

Maximum change to a log proposal scale in one warmup transition.

SGLDConfig dataclass

Bases: BaseMCMCConfig

Stochastic Gradient Langevin Dynamics (SGLD) configuration.

SGLD performs first-order Langevin updates using noisy gradients and injected Gaussian noise. In this implementation, gradients are computed on the full dataset (no minibatching), so the method behaves as full-batch Langevin dynamics with an annealed step schedule.

Attributes:

Name Type Description
step_size float

Base learning rate used in the SGLD schedule. This should generally be small.

schedule_power float

Power in the polynomial decay schedule \(\epsilon_t = \text{step_size} \cdot t^{-\text{schedule_power}}\). Values in (0.5, 1.0] are common for asymptotic convergence.

MALAConfig dataclass

Bases: BaseMCMCConfig

Metropolis-Adjusted Langevin Algorithm (MALA) configuration.

Attributes:

Name Type Description
step_size float

Proposal step size used by blackjax.mala.