MCMC

Markov-chain Monte Carlo: the HMC and NUTS kernels, the MCMC runner, and the MCMCResult summary.

The runner targets a MonadicProgram. Models that declare every latent as a sample site need no conversion. For models with nn.Parameters or intermediate latent sites, the bayesian_lift_parameters lift returns a 3-tuple (MonadicProgram, torch.Tensor, dict[str, torch.Tensor]), whose first element is the matching MonadicProgram. Callers pass that first element to the runner.

quivers.inference re-exports the kernels, the runner, and its result; quivers.inference.mcmc also exports the warmup adaptation primitives and the type aliases that parameterize the kernels and the runner.

mcmc

MCMC kernels, adaptation primitives, and the chain driver.

Public surface (the kernels and driver are also re-exported from quivers.inference):

• MCMCKernel: ABC for Markov kernels on the flat unconstrained latent vector, with KernelState as its per-chain state and PotentialFn as the potential it evaluates.
• HMCKernel: Hamiltonian Monte Carlo with leapfrog integration, dual-averaging step-size adaptation, and Welford mass-matrix adaptation, parameterized by a MassMatrixKind.
• NUTSKernel: No-U-Turn Sampler with multinomial sampling.
• DualAveraging, WelfordCovariance, and find_reasonable_step_size: the warmup adaptation primitives the kernels share.
• MCMC: chain orchestrator with warmup, parallel chains, an InitStrategy, and posterior diagnostics (split-:math:\hat R, effective sample size).
• MCMCResult: posterior samples and per-chain diagnostics.

InitStrategy module-attribute

InitStrategy = Literal['prior', 'zero', 'guide']

How MCMC chooses each chain's initial point in the unconstrained latent space: "prior" draws a small seeded standard-normal perturbation of the origin, "zero" starts at the origin, and "guide" maps one draw of a fitted guide into the unconstrained space.

MassMatrixKind module-attribute

MassMatrixKind = Literal['identity', 'diagonal', 'dense']

The mass-matrix structure an HMC or NUTS kernel adapts during warmup: "identity" keeps the unit metric, "diagonal" adapts per-coordinate variances, and "dense" adapts a full covariance.

DualAveraging

DualAveraging(initial_step_size: float, target_accept: float = 0.8, gamma: float = 0.05, t0: float = 10.0, kappa: float = 0.75)

Nesterov dual averaging for adaptive step-size selection.

Maintains a running smoothed estimate of the log-step-size that drives the chain's average acceptance probability toward the user-supplied target_accept.

Algorithm parameters follow Hoffman-Gelman 2014 (Algorithm 6): :math:\gamma = 0.05, :math:t_0 = 10, :math:\kappa = 0.75. These are the defaults used by Stan.

PARAMETER	DESCRIPTION
initial_step_size	Starting step size. The first proposal uses log(10 * ε₀) as the shrinkage anchor :math:\mu per Algorithm 6. TYPE: float
target_accept	Target Metropolis acceptance probability. Default 0.8 (NUTS's standard), somewhat below the optimal 0.65 for HMC because NUTS averages acceptance over the tree. TYPE: float DEFAULT: 0.8
gamma	Adaptation regularisation. Default 0.05. TYPE: float DEFAULT: 0.05
t0	Stabilisation offset. Default 10.0. TYPE: float DEFAULT: 10.0
kappa	Decay exponent for the smoothed step size. Default 0.75. TYPE: float DEFAULT: 0.75

Source code in src/quivers/inference/mcmc/adapt.py

def __init__(
    self,
    initial_step_size: float,
    target_accept: float = 0.8,
    gamma: float = 0.05,
    t0: float = 10.0,
    kappa: float = 0.75,
) -> None:
    if initial_step_size <= 0:
        raise ValueError(
            f"DualAveraging: initial_step_size must be > 0, got {initial_step_size}"
        )
    if not 0.0 < target_accept < 1.0:
        raise ValueError(
            f"DualAveraging: target_accept must be in (0, 1), got {target_accept}"
        )
    self._mu = math.log(10.0 * initial_step_size)
    self._target_accept = target_accept
    self._gamma = gamma
    self._t0 = t0
    self._kappa = kappa
    self._log_step = math.log(initial_step_size)
    self._log_step_smoothed = 0.0
    self._h_bar = 0.0
    self._step_count = 0

step_size

step_size() -> float

Current step size used by the next leapfrog proposal.

Source code in src/quivers/inference/mcmc/adapt.py

def step_size(self) -> float:
    """Current step size used by the *next* leapfrog proposal."""
    return math.exp(self._log_step)

smoothed_step_size

smoothed_step_size() -> float

Smoothed step size frozen after warmup ends.

Source code in src/quivers/inference/mcmc/adapt.py

def smoothed_step_size(self) -> float:
    """Smoothed step size frozen after warmup ends."""
    return math.exp(self._log_step_smoothed)

update

update(acceptance_prob: float) -> None

Update the step size given the latest Metropolis acceptance probability.

Source code in src/quivers/inference/mcmc/adapt.py

def update(self, acceptance_prob: float) -> None:
    """Update the step size given the latest Metropolis
    acceptance probability."""
    if not math.isfinite(acceptance_prob):
        acceptance_prob = 0.0
    acceptance_prob = max(0.0, min(1.0, acceptance_prob))
    self._step_count += 1
    m = float(self._step_count)
    eta_h = 1.0 / (m + self._t0)
    self._h_bar = (1.0 - eta_h) * self._h_bar + eta_h * (
        self._target_accept - acceptance_prob
    )
    self._log_step = self._mu - math.sqrt(m) / self._gamma * self._h_bar
    eta_step = m ** (-self._kappa)
    self._log_step_smoothed = (
        eta_step * self._log_step + (1.0 - eta_step) * self._log_step_smoothed
    )

WelfordCovariance

WelfordCovariance(dim: int, regularise: bool = True, diagonal: bool = False)

Welford-style online covariance accumulator.

Tracks the running mean and (co)variance of a stream of D-dimensional vectors. Supports both diagonal and dense forms: diagonal is sufficient for ill-scaled axis-aligned posteriors, dense needed when off-axis correlations matter.

The dense form follows the standard pairwise update:

.. math::

M_n = M_{n-1} + (x_n - \bar{x}_{n-1})(x_n - \bar{x}_n)^\top

with :math:\bar{x}_n = \bar{x}_{n-1} + (x_n - \bar{x}_{n-1}) / n.

PARAMETER	DESCRIPTION
dim	Vector dimension. TYPE: int
regularise	Apply Stan's regularisation :math:(n / (n + 5)) \cdot \Sigma + 0.001 \cdot (5 / (n + 5)) \cdot I to the final covariance. Default True; mirrors Stan's warmup-end mass-matrix construction. TYPE: bool DEFAULT: True
diagonal	Track only the diagonal of the covariance (vector, not matrix). Default False. TYPE: bool DEFAULT: False

Source code in src/quivers/inference/mcmc/adapt.py

def __init__(
    self,
    dim: int,
    regularise: bool = True,
    diagonal: bool = False,
) -> None:
    if dim < 1:
        raise ValueError(f"WelfordCovariance: dim must be >= 1, got {dim}")
    self._dim = dim
    self._regularise = regularise
    self._diagonal = diagonal
    self._n = 0
    self._mean = torch.zeros(dim)
    self._m2: torch.Tensor
    if diagonal:
        self._m2 = torch.zeros(dim)
    else:
        self._m2 = torch.zeros(dim, dim)

update

update(x: Tensor) -> None

Fold one vector into the running statistics.

Source code in src/quivers/inference/mcmc/adapt.py

def update(self, x: torch.Tensor) -> None:
    """Fold one vector into the running statistics."""
    if x.shape != (self._dim,):
        raise ValueError(
            f"WelfordCovariance.update: expected shape "
            f"({self._dim},); got {tuple(x.shape)}"
        )
    self._n += 1
    delta = x - self._mean
    self._mean = self._mean + delta / float(self._n)
    delta_post = x - self._mean
    if self._diagonal:
        self._m2 = self._m2 + delta * delta_post
    else:
        self._m2 = self._m2 + delta.unsqueeze(-1) @ delta_post.unsqueeze(0)

covariance

covariance() -> Tensor

Return the empirical covariance (regularised if configured).

Source code in src/quivers/inference/mcmc/adapt.py

def covariance(self) -> torch.Tensor:
    """Return the empirical covariance (regularised if
    configured)."""
    if self._n < 2:
        # Insufficient samples; return identity so the kernel
        # behaves as before adaptation kicked in.
        if self._diagonal:
            return torch.ones(self._dim)
        return torch.eye(self._dim)
    cov = self._m2 / float(self._n - 1)
    if self._regularise:
        n = float(self._n)
        shrink = n / (n + 5.0)
        jitter = 1e-3 * (5.0 / (n + 5.0))
        if self._diagonal:
            cov = shrink * cov + jitter * torch.ones(self._dim)
        else:
            cov = shrink * cov + jitter * torch.eye(self._dim)
    return cov

reset

reset() -> None

Discard accumulated statistics.

Source code in src/quivers/inference/mcmc/adapt.py

def reset(self) -> None:
    """Discard accumulated statistics."""
    self._n = 0
    self._mean = torch.zeros(self._dim)
    if self._diagonal:
        self._m2 = torch.zeros(self._dim)
    else:
        self._m2 = torch.zeros(self._dim, self._dim)

MCMC

MCMC(kernel: MCMCKernel, num_warmup: int, num_samples: int, num_chains: int = 4, init_strategy: InitStrategy = 'prior', *, seed: int | None = None, generator: Generator | None = None)

MCMC chain runner.

PARAMETER	DESCRIPTION
kernel	Markov kernel (e.g. HMCKernel, NUTSKernel). TYPE: MCMCKernel
num_warmup	Number of adaptation steps. The kernel's adaptation machinery (dual averaging, Welford covariance) runs over this prefix. TYPE: int
num_samples	Post-warmup samples per chain. TYPE: int
num_chains	Independent chains. Default 4 (Stan / NumPyro default). TYPE: int DEFAULT: 4
init_strategy	How to pick each chain's initial position. TYPE: ('prior', 'zero', 'guide') DEFAULT: "prior"
seed	Base seed used to derive an independent local generator for each chain. Repeated calls to run with the same seed are reproducible and do not consume Torch's global RNG. TYPE: int DEFAULT: None
generator	Caller-owned CPU generator used to draw one base seed per run. The generator advances, while Torch's global RNG does not. Mutually exclusive with seed. TYPE: Generator DEFAULT: None

Source code in src/quivers/inference/mcmc/driver.py

def __init__(
    self,
    kernel: MCMCKernel,
    num_warmup: int,
    num_samples: int,
    num_chains: int = 4,
    init_strategy: InitStrategy = "prior",
    *,
    seed: int | None = None,
    generator: torch.Generator | None = None,
) -> None:
    if num_warmup < 0:
        raise ValueError(f"MCMC: num_warmup must be >= 0, got {num_warmup}")
    if num_samples < 1:
        raise ValueError(f"MCMC: num_samples must be >= 1, got {num_samples}")
    if num_chains < 1:
        raise ValueError(f"MCMC: num_chains must be >= 1, got {num_chains}")
    if seed is not None and generator is not None:
        raise ValueError("MCMC: seed and generator are mutually exclusive")
    if seed is not None and (not isinstance(seed, int) or isinstance(seed, bool)):
        raise TypeError(f"MCMC: seed must be an int, got {type(seed).__name__}")
    if generator is not None:
        if not isinstance(generator, torch.Generator):
            raise TypeError(
                "MCMC: generator must be a torch.Generator, "
                f"got {type(generator).__name__}"
            )
        if generator.device.type != "cpu":
            raise ValueError(
                "MCMC: generator must be a CPU generator because MCMC "
                "currently runs chain state on CPU"
            )
    self.kernel = kernel
    self.num_warmup = num_warmup
    self.num_samples = num_samples
    self.num_chains = num_chains
    self.init_strategy = init_strategy
    self.seed = seed
    self.generator = generator

run

run(model: MonadicProgram, x: Tensor, observations: dict[str, Tensor], guide: Guide | None = None) -> MCMCResult

Run the configured kernel for num_chains chains of num_warmup + num_samples steps each.

Source code in src/quivers/inference/mcmc/driver.py

def run(
    self,
    model: MonadicProgram,
    x: torch.Tensor,
    observations: dict[str, torch.Tensor],
    guide: Guide | None = None,
) -> MCMCResult:
    """Run the configured kernel for ``num_chains`` chains of
    ``num_warmup + num_samples`` steps each."""
    observed_names = set(observations.keys())
    registry = LatentRegistry.from_model(model, observed_names)
    potential = PotentialFn(model, registry, x, observations)

    site_shapes: dict[str, tuple[int, ...]] = {
        site.name: site.constrained_shape or (1,)
        for site in registry.sites.values()
    }
    per_chain_samples: dict[str, list[torch.Tensor]] = {n: [] for n in site_shapes}
    per_chain_log_density = torch.empty(self.num_chains, self.num_samples)
    per_chain_accept = torch.empty(self.num_chains)
    per_chain_divergences = torch.empty(self.num_chains, dtype=torch.int64)
    run_seed = self._run_seed()

    for chain in range(self.num_chains):
        chain_generator = self._chain_generator(run_seed, chain)
        init_pos = self._initial_position(registry, guide, chain_generator)
        if not torch.isfinite(init_pos).all():
            raise ValueError(
                f"MCMC.run: chain {chain} initial_position must contain "
                "only finite values"
            )
        self.kernel._set_generator(chain_generator)
        state = self.kernel.init(registry, model, x, observations, init_pos)
        # Warmup.
        if self.num_warmup > 0:
            self.kernel.start_adaptation()
            for _ in range(self.num_warmup):
                state = self.kernel.step(state, potential)
            self.kernel.stop_adaptation()
        # Reset accept-count so the reported rate excludes warmup.
        sampling_state = KernelState(
            position=state.position,
            log_density=state.log_density,
            grad_log_density=state.grad_log_density,
            step_count=state.step_count,
            accept_count=0,
            diverged=False,
            extras=state.extras,
        )
        chain_samples: dict[str, list[torch.Tensor]] = {n: [] for n in site_shapes}
        sampling_divergences = 0
        for s in range(self.num_samples):
            sampling_state = self.kernel.step(sampling_state, potential)
            if sampling_state.diverged:
                sampling_divergences += 1
            draws = self._to_constrained(registry, sampling_state.position)
            for n, v in draws.items():
                chain_samples[n].append(v)
            per_chain_log_density[chain, s] = sampling_state.log_density
        per_chain_accept[chain] = sampling_state.accept_count / float(
            self.num_samples
        )
        per_chain_divergences[chain] = sampling_divergences
        for n, draws_list in chain_samples.items():
            stacked = torch.stack(draws_list, dim=0)
            per_chain_samples[n].append(stacked)

    # Stack per-chain → (num_chains, num_samples, *site_shape).
    samples: dict[str, torch.Tensor] = {}
    for n, chain_draws in per_chain_samples.items():
        samples[n] = torch.stack(chain_draws, dim=0)

    r_hat, ess = _diagnostics(samples)
    return MCMCResult(
        samples=samples,
        log_densities=per_chain_log_density,
        acceptance_rates=per_chain_accept,
        divergence_counts=per_chain_divergences,
        r_hat=r_hat,
        ess=ess,
        num_warmup=self.num_warmup,
        num_samples=self.num_samples,
    )

MCMCResult dataclass

MCMCResult(samples: dict[str, Tensor], log_densities: Tensor, acceptance_rates: Tensor, divergence_counts: Tensor, r_hat: dict[str, Tensor], ess: dict[str, Tensor], num_warmup: int, num_samples: int)

Posterior samples and per-chain diagnostics.

ATTRIBUTE	DESCRIPTION
samples	Per-site posterior draws on the constrained support. Shape (num_chains, num_samples, *site_shape). TYPE: dict[str, Tensor]
log_densities	Unconstrained-space log-density (Jacobian-corrected) at every posterior draw. Shape (num_chains, num_samples). TYPE: Tensor
acceptance_rates	Per-chain post-warmup acceptance rate. Shape (num_chains,). TYPE: Tensor
divergence_counts	Per-chain post-warmup divergence count. Shape (num_chains,). TYPE: Tensor
r_hat	Per-site rank-normalized folded split-:math:\hat R. Each coordinate reports the conservative maximum of the rank-normalized and folded statistics. Each site's tensor has the site's shape (one scalar per coordinate). TYPE: dict[str, Tensor]
ess	Per-site rank-normalized bulk effective sample size pooled across chains. Same shape convention as r_hat. TYPE: dict[str, Tensor]
num_warmup	TYPE: int
num_samples	TYPE: int

HMCKernel

HMCKernel(step_size: float = 0.1, num_steps: int = 10, mass_matrix: MassMatrixKind = 'identity', target_accept: float = 0.65, divergence_threshold: float = 1000.0, adapt_step_size: bool = True, adapt_mass_matrix: bool = True)

Bases: MCMCKernel

Hamiltonian Monte Carlo kernel with fixed trajectory length.

PARAMETER	DESCRIPTION
step_size	Leapfrog step size. Adapted during warmup when adapt_step_size is true. TYPE: float DEFAULT: 0.1
num_steps	Leapfrog steps per proposal. TYPE: int DEFAULT: 10
mass_matrix	Mass-matrix shape. Diagonal / dense are adapted during warmup from the empirical covariance of warmup samples. TYPE: ('identity', 'diagonal', 'dense') DEFAULT: "identity"
target_accept	Target Metropolis acceptance for dual averaging. Default 0.65 (Beskos et al.'s HMC-optimal acceptance for product-form targets). TYPE: float DEFAULT: 0.65
divergence_threshold	Energy-error threshold for marking a proposal as divergent. Divergent steps still respect Metropolis correctness but are reported separately so the user can spot pathological regions. TYPE: float DEFAULT: 1000.0

Source code in src/quivers/inference/mcmc/hmc.py

def __init__(
    self,
    step_size: float = 0.1,
    num_steps: int = 10,
    mass_matrix: MassMatrixKind = "identity",
    target_accept: float = 0.65,
    divergence_threshold: float = 1000.0,
    adapt_step_size: bool = True,
    adapt_mass_matrix: bool = True,
) -> None:
    if step_size <= 0:
        raise ValueError(f"HMCKernel: step_size must be > 0, got {step_size}")
    if num_steps < 1:
        raise ValueError(f"HMCKernel: num_steps must be >= 1, got {num_steps}")
    if not 0.0 < target_accept < 1.0:
        raise ValueError(
            f"HMCKernel: target_accept must be in (0, 1), got {target_accept}"
        )
    self._step_size = step_size
    self._num_steps = num_steps
    self._mass_kind = mass_matrix
    self._target_accept = target_accept
    self._divergence_threshold = divergence_threshold
    self._adapt_step_size = adapt_step_size
    self._adapt_mass_matrix = adapt_mass_matrix and mass_matrix != "identity"
    self._mass: _MassMatrix | None = None
    self._dual_avg: DualAveraging | None = None
    self._welford: WelfordCovariance | None = None

NUTSKernel

NUTSKernel(step_size: float = 0.1, max_tree_depth: int = 10, mass_matrix: MassMatrixKind = 'diagonal', target_accept: float = 0.8, divergence_threshold: float = 1000.0, adapt_step_size: bool = True, adapt_mass_matrix: bool = True)

Bases: MCMCKernel

No-U-Turn Sampler with multinomial sampling and the standard U-turn termination (Hoffman-Gelman 2014 algorithms 3 + 6, Betancourt 2017's generalized slice variant for multinomial sampling).

PARAMETER	DESCRIPTION
step_size	Initial leapfrog step size; adapted via dual averaging. TYPE: float DEFAULT: 0.1
max_tree_depth	Maximum tree doubling depth. Default 10 (Stan default). TYPE: int DEFAULT: 10
target_accept	Target tree-averaged Metropolis acceptance for dual averaging. Default 0.8. TYPE: float DEFAULT: 0.8
divergence_threshold	Energy-error threshold above which a leapfrog substep is marked divergent and terminates the tree on its branch. TYPE: float DEFAULT: 1000.0

Source code in src/quivers/inference/mcmc/hmc.py

def __init__(
    self,
    step_size: float = 0.1,
    max_tree_depth: int = 10,
    mass_matrix: MassMatrixKind = "diagonal",
    target_accept: float = 0.8,
    divergence_threshold: float = 1000.0,
    adapt_step_size: bool = True,
    adapt_mass_matrix: bool = True,
) -> None:
    if step_size <= 0:
        raise ValueError(f"NUTSKernel: step_size must be > 0, got {step_size}")
    if max_tree_depth < 1:
        raise ValueError(
            f"NUTSKernel: max_tree_depth must be >= 1, got {max_tree_depth}"
        )
    if not 0.0 < target_accept < 1.0:
        raise ValueError(
            f"NUTSKernel: target_accept must be in (0, 1), got {target_accept}"
        )
    self._step_size = step_size
    self._max_depth = max_tree_depth
    self._mass_kind = mass_matrix
    self._target_accept = target_accept
    self._divergence_threshold = divergence_threshold
    self._adapt_step_size = adapt_step_size
    self._adapt_mass_matrix = adapt_mass_matrix and mass_matrix != "identity"
    self._mass: _MassMatrix | None = None
    self._dual_avg: DualAveraging | None = None
    self._welford: WelfordCovariance | None = None

KernelState dataclass

KernelState(position: Tensor, log_density: Tensor, grad_log_density: Tensor, step_count: int = 0, accept_count: int = 0, diverged: bool = False, extras: dict = dict())

Mutable container for one chain's MCMC state.

ATTRIBUTE	DESCRIPTION
position	Current flat unconstrained latent vector. Shape (D,) for one chain, (num_chains, D) for parallel chains. TYPE: Tensor
log_density	Unconstrained-space log-density (Jacobian-corrected log-joint) at position. Shape () or (num_chains,). TYPE: Tensor
grad_log_density	Gradient of log_density with respect to position, same shape as position. Cached so leapfrog re-uses the last-evaluated gradient at the proposal endpoint. TYPE: Tensor
step_count	Number of step calls so far (counts both warmup and post-warmup steps). TYPE: int
accept_count	Number of proposals accepted across the chain. Useful for reporting acceptance rate. TYPE: int
diverged	Whether the most recent step's energy error exceeded the kernel's divergence threshold. Reset by each kernel as appropriate. TYPE: bool
extras	Per-kernel additional state (e.g. NUTS tree depth, HMC step-size adaptation cumulants). TYPE: dict

MCMCKernel

Bases: ABC

Abstract Markov kernel on the flat unconstrained latent vector.

Concrete subclasses implement init and step. Adaptation phases (warmup) typically mutate kernel-internal state (step size, mass matrix) and freeze it for the sampling phase; the kernel's is_adapting flag tracks that.

init abstractmethod

init(registry: LatentRegistry, model: MonadicProgram, x: Tensor, observations: dict[str, Tensor], initial_position: Tensor) -> KernelState

Build the starting KernelState from the supplied initial flat unconstrained vector. The initial gradient is evaluated here so step can re-use it.

Source code in src/quivers/inference/mcmc/kernel.py

@abstractmethod
def init(
    self,
    registry: LatentRegistry,
    model: MonadicProgram,
    x: torch.Tensor,
    observations: dict[str, torch.Tensor],
    initial_position: torch.Tensor,
) -> KernelState:
    """Build the starting `KernelState` from the supplied
    initial flat unconstrained vector. The initial gradient is
    evaluated here so `step` can re-use it."""

step abstractmethod

step(state: KernelState, potential: PotentialFn) -> KernelState

Advance the chain one Metropolis step. The potential function is constructed once per MCMC.run and re-used across every step / chain.

Source code in src/quivers/inference/mcmc/kernel.py

@abstractmethod
def step(
    self,
    state: KernelState,
    potential: PotentialFn,
) -> KernelState:
    """Advance the chain one Metropolis step. The potential
    function is constructed once per `MCMC.run` and
    re-used across every step / chain."""

start_adaptation

start_adaptation() -> None

Enter the adaptation (warmup) phase.

Source code in src/quivers/inference/mcmc/kernel.py

def start_adaptation(self) -> None:
    """Enter the adaptation (warmup) phase."""
    self.is_adapting = True

stop_adaptation

stop_adaptation() -> None

Freeze the kernel's adapted parameters.

Source code in src/quivers/inference/mcmc/kernel.py

def stop_adaptation(self) -> None:
    """Freeze the kernel's adapted parameters."""
    self.is_adapting = False

PotentialFn

PotentialFn(model: MonadicProgram, registry: LatentRegistry, x: Tensor, observations: dict[str, Tensor])

Callable that maps a flat unconstrained position to the unconstrained-space negative log-density and its gradient.

HMC and NUTS need both the potential :math:U(z) = -\log \tilde{p}(z) (where :math:\tilde{p}(z) = p(T(z), y) \cdot |\det J_T(z)| is the Jacobian-corrected unconstrained-space joint) and its gradient :math:\nabla U(z). The two are computed in a single autograd pass and cached on the kernel state.

PARAMETER	DESCRIPTION
model	Generative model. TYPE: MonadicProgram
registry	Latent-site registry for model. TYPE: LatentRegistry
x	Model input. Shape (batch, ...). TYPE: Tensor
observations	Observed-site values and host data. TYPE: dict[str, Tensor]

Source code in src/quivers/inference/mcmc/kernel.py

def __init__(
    self,
    model: MonadicProgram,
    registry: LatentRegistry,
    x: torch.Tensor,
    observations: dict[str, torch.Tensor],
) -> None:
    self._model = model
    self._registry = registry
    self._x = x
    self._observations = observations

log_density

log_density(z: Tensor) -> Tensor

Unconstrained-space log-density (Jacobian-corrected).

Trajectories that wander to the edge of a constrained support can produce values that fall outside torch.distributions' validation envelope (e.g. exact zeros against a strictly-positive support after a long leapfrog stride). Rather than letting the resulting ValueError propagate and kill the chain, this method returns -inf for those positions; the kernel reads non-finite log-densities as divergent transitions and rejects them in the Metropolis step.

Source code in src/quivers/inference/mcmc/kernel.py

def log_density(self, z: torch.Tensor) -> torch.Tensor:
    """Unconstrained-space log-density (Jacobian-corrected).

    Trajectories that wander to the edge of a constrained
    support can produce values that fall outside
    `torch.distributions`' validation envelope (e.g. exact
    zeros against a strictly-positive support after a long
    leapfrog stride). Rather than letting the resulting
    ``ValueError`` propagate and kill the chain, this method
    returns ``-inf`` for those positions; the kernel reads
    non-finite log-densities as divergent transitions and
    rejects them in the Metropolis step.
    """
    try:
        per_site_unc = self._registry.unflatten_unconstrained(z)
        constrained: dict[str, torch.Tensor] = {}
        log_det_total = torch.zeros((), device=z.device, dtype=z.dtype)
        for site in self._registry.sites.values():
            z_site = per_site_unc[site.name]
            if not site.is_plate:
                z_site = z_site.unsqueeze(0)
            v = site.bijector(z_site)
            log_det_total = log_det_total + (
                site.bijector.log_abs_det_jacobian(z_site, v).sum()
            )
            if site.constrained_dim == 1 and v.dim() >= 1 and v.shape[-1] == 1:
                v = v.squeeze(-1)
            constrained[site.name] = v
        log_p = self._model.log_joint(
            self._x, {**constrained, **self._observations}
        )
        result = log_p.sum() + log_det_total
    except ValueError:
        return torch.tensor(float("-inf"), device=z.device, dtype=z.dtype)
    if not torch.isfinite(result):
        return torch.tensor(float("-inf"), device=z.device, dtype=z.dtype)
    return result

value_and_grad

value_and_grad(z: Tensor) -> tuple[Tensor, Tensor]

Return (log_density, grad_log_density) for z.

z is expected to be a detached tensor; we make a fresh leaf with requires_grad=True so gradient propagation doesn't leak into the kernel's accumulated state.

For divergent positions (where the log-density is -inf), returns a zero gradient: the kernel rejects the trajectory in the Metropolis step anyway, and a zero gradient keeps the leapfrog integrator from producing NaN downstream.

Source code in src/quivers/inference/mcmc/kernel.py

def value_and_grad(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Return ``(log_density, grad_log_density)`` for ``z``.

    ``z`` is expected to be a detached tensor; we make a fresh
    leaf with ``requires_grad=True`` so gradient propagation
    doesn't leak into the kernel's accumulated state.

    For divergent positions (where the log-density is
    ``-inf``), returns a zero gradient: the kernel rejects
    the trajectory in the Metropolis step anyway, and a zero
    gradient keeps the leapfrog integrator from producing NaN
    downstream.
    """
    z_leaf = z.detach().clone().requires_grad_(True)
    try:
        ld = self.log_density(z_leaf)
        if not torch.isfinite(ld):
            return ld.detach(), torch.zeros_like(z)
        grad = torch.autograd.grad(
            ld, z_leaf, create_graph=False, allow_unused=False
        )[0]
        if grad is None or not torch.isfinite(grad).all():
            return ld.detach(), torch.zeros_like(z)
        return ld.detach(), grad.detach()
    except ValueError, RuntimeError:
        return (
            torch.tensor(float("-inf"), device=z.device, dtype=z.dtype),
            torch.zeros_like(z),
        )

find_reasonable_step_size

find_reasonable_step_size(log_density: Callable[[Tensor], Tensor], grad_log_density: Callable[[Tensor], Tensor], z0: Tensor, initial_step: float = 1.0, target_log_accept: float = log(0.5)) -> float

Heuristic-by-doubling initial step size (Hoffman-Gelman 2014 Algorithm 4).

Starting from initial_step, doubles or halves the step size until a single leapfrog step's Metropolis log-acceptance crosses target_log_accept (default :math:\log 0.5). Returns the step size at the crossing.

Source code in src/quivers/inference/mcmc/adapt.py

def find_reasonable_step_size(
    log_density: Callable[[torch.Tensor], torch.Tensor],
    grad_log_density: Callable[[torch.Tensor], torch.Tensor],
    z0: torch.Tensor,
    initial_step: float = 1.0,
    target_log_accept: float = math.log(0.5),
) -> float:
    """Heuristic-by-doubling initial step size (Hoffman-Gelman 2014
    Algorithm 4).

    Starting from ``initial_step``, doubles or halves the step size
    until a single leapfrog step's Metropolis log-acceptance crosses
    ``target_log_accept`` (default :math:`\\log 0.5`). Returns the
    step size at the crossing.
    """
    eps = initial_step
    z = z0.clone()
    ld0 = log_density(z)
    g0 = grad_log_density(z)
    p = torch.randn_like(z)
    h0 = ld0 - 0.5 * (p * p).sum()
    # One leapfrog step at eps.
    z1 = z + eps * p + 0.5 * eps * eps * g0
    g1 = grad_log_density(z1)
    p1 = p + 0.5 * eps * (g0 + g1)
    ld1 = log_density(z1)
    h1 = ld1 - 0.5 * (p1 * p1).sum()
    log_accept = float(h1 - h0)
    if not math.isfinite(log_accept):
        log_accept = -float("inf")
    direction = 1 if log_accept > target_log_accept else -1
    for _ in range(100):
        eps = eps * (2.0**direction)
        if eps < 1e-10 or eps > 1e10:
            break
        z1 = z + eps * p + 0.5 * eps * eps * g0
        g1 = grad_log_density(z1)
        p1 = p + 0.5 * eps * (g0 + g1)
        ld1 = log_density(z1)
        h1 = ld1 - 0.5 * (p1 * p1).sum()
        new_log_accept = float(h1 - h0)
        if not math.isfinite(new_log_accept):
            new_log_accept = -float("inf")
        crossed = (direction == 1 and new_log_accept < target_log_accept) or (
            direction == -1 and new_log_accept > target_log_accept
        )
        log_accept = new_log_accept
        if crossed:
            break
    return eps

Variational warmup

WarmupThenHMC fits a guide with SVI and starts the chains from it.

warmup

Run variational warmup before HMC or NUTS.

The orchestrator fits a guide with SVI, initializes the Markov chain from the guide, adapts the mass matrix from its covariance, and runs the configured MCMC driver.

WarmupThenHMC

WarmupThenHMC(guide: Guide, kernel: MCMCKernel, svi_steps: int, mcmc_warmup: int, mcmc_samples: int, num_chains: int = 2, svi_lr: float = 0.01, objective: Objective | None = None)

Composite sampler: SVI warmup followed by HMC / NUTS chain.

PARAMETER	DESCRIPTION
guide	Variational guide trained in the SVI warmup phase. Any Guide subclass; AutoMultivariateNormalGuide is the canonical pick because its learned Cholesky factor directly seeds the dense mass matrix of HMC's adaptation. TYPE: Guide
kernel	Markov kernel for the sampling phase (typically HMCKernel or NUTSKernel). TYPE: MCMCKernel
svi_steps	Number of SVI steps in the warmup phase. TYPE: int
svi_lr	Learning rate for the SVI optimizer. Default 1e-2. TYPE: float DEFAULT: 0.01
mcmc_warmup	MCMC kernel's own warmup (on top of the SVI warmup; HMC still benefits from a few step-size-adaptation iterations because the SVI-mean's typical-set position differs from the maximum-density point). TYPE: int
mcmc_samples	Number of post-warmup MCMC samples per chain. TYPE: int
num_chains	Number of MCMC chains. TYPE: int DEFAULT: 2
objective	SVI objective. Default ELBO. TYPE: Objective DEFAULT: None

Source code in src/quivers/inference/warmup.py

def __init__(
    self,
    guide: Guide,
    kernel: MCMCKernel,
    svi_steps: int,
    mcmc_warmup: int,
    mcmc_samples: int,
    num_chains: int = 2,
    svi_lr: float = 1e-2,
    objective: Objective | None = None,
) -> None:
    if svi_steps < 1:
        raise ValueError(f"WarmupThenHMC: svi_steps must be >= 1, got {svi_steps}")
    if mcmc_warmup < 0:
        raise ValueError(
            f"WarmupThenHMC: mcmc_warmup must be >= 0, got {mcmc_warmup}"
        )
    if mcmc_samples < 1:
        raise ValueError(
            f"WarmupThenHMC: mcmc_samples must be >= 1, got {mcmc_samples}"
        )
    self.guide = guide
    self.kernel = kernel
    self.svi_steps = svi_steps
    self.svi_lr = svi_lr
    self.mcmc_warmup = mcmc_warmup
    self.mcmc_samples = mcmc_samples
    self.num_chains = num_chains
    self.objective = objective if objective is not None else ELBO()

fit_guide

fit_guide(model: MonadicProgram, x: Tensor, observations: dict[str, Tensor]) -> list[float]

Run the SVI warmup phase and return the per-step losses.

Source code in src/quivers/inference/warmup.py

def fit_guide(
    self,
    model: MonadicProgram,
    x: torch.Tensor,
    observations: dict[str, torch.Tensor],
) -> list[float]:
    """Run the SVI warmup phase and return the per-step losses."""
    optim = torch.optim.Adam(
        list(model.parameters()) + list(self.guide.parameters()),
        lr=self.svi_lr,
    )
    svi = SVI(model, self.guide, optim, self.objective)
    losses: list[float] = []
    for _ in range(self.svi_steps):
        losses.append(svi.step(x, observations))
    return losses

run_mcmc

run_mcmc(model: MonadicProgram, x: Tensor, observations: dict[str, Tensor]) -> MCMCResult

Run the post-warmup MCMC phase. Chains are seeded from the fitted guide via MCMC.run's init_strategy='guide'.

Source code in src/quivers/inference/warmup.py

def run_mcmc(
    self,
    model: MonadicProgram,
    x: torch.Tensor,
    observations: dict[str, torch.Tensor],
) -> MCMCResult:
    """Run the post-warmup MCMC phase. Chains are seeded from
    the fitted guide via `MCMC.run`'s ``init_strategy='guide'``."""
    driver = MCMC(
        kernel=self.kernel,
        num_warmup=self.mcmc_warmup,
        num_samples=self.mcmc_samples,
        num_chains=self.num_chains,
        init_strategy="guide",
    )
    return driver.run(model, x, observations, guide=self.guide)

run

run(model: MonadicProgram, x: Tensor, observations: dict[str, Tensor]) -> tuple[list[float], MCMCResult]

Convenience: SVI warmup, then MCMC sampling.

Returns (svi_losses, mcmc_result).

Source code in src/quivers/inference/warmup.py

def run(
    self,
    model: MonadicProgram,
    x: torch.Tensor,
    observations: dict[str, torch.Tensor],
) -> tuple[list[float], MCMCResult]:
    """Convenience: SVI warmup, then MCMC sampling.

    Returns ``(svi_losses, mcmc_result)``.
    """
    losses = self.fit_guide(model, x, observations)
    result = self.run_mcmc(model, x, observations)
    return losses, result