Implementation:Pyro ppl Pyro HMC Kernel

Metadata

Field	Value
Page Type	Implementation (API Doc)
Knowledge Sources	Repo (Pyro), Paper (MCMC Using Hamiltonian Dynamics)
Domains	MCMC, Bayesian_Inference
Last Updated	2026-02-09 12:00 GMT

Overview

Concrete MCMC kernel implementing the Hamiltonian Monte Carlo algorithm in Pyro, providing gradient-based sampling with leapfrog integration, step size adaptation, and mass matrix adaptation.

Description

HMC is a subclass of MCMCKernel that implements the standard Hamiltonian Monte Carlo algorithm. At each iteration, HMC:

Samples new momentum variables from a multivariate normal distribution parameterized by the mass matrix.
Simulates Hamiltonian dynamics using the leapfrog integrator for a fixed number of steps (or a fixed trajectory length).
Applies a Metropolis accept/reject correction to account for discretization error.

The kernel supports the following adaptation mechanisms during warmup:

Step size adaptation: Via dual averaging, targeting a specified acceptance probability.
Mass matrix adaptation: Via the Welford online algorithm, estimating either diagonal or full covariance.
Automatic transforms: Constrained parameters (e.g., positive-valued, bounded) are automatically reparameterized to unconstrained space using appropriate bijective transforms from Pyro's constraint registry.

The trajectory length can be specified either as num_steps (number of leapfrog steps) or as trajectory_length (total integration time, with the number of steps computed as trajectory_length / step_size). If neither is provided, trajectory_length defaults to 2 * pi.

Code Reference

Source Location

Pyro repo, file: pyro/infer/mcmc/hmc.py, lines L21-453.

Class Hierarchy

HMC inherits from MCMCKernel (defined in pyro/infer/mcmc/mcmc_kernel.py).

Signature

class HMC(MCMCKernel):
    def __init__(
        self,
        model=None,
        potential_fn=None,
        step_size=1,
        trajectory_length=None,
        num_steps=None,
        adapt_step_size=True,
        adapt_mass_matrix=True,
        full_mass=False,
        transforms=None,
        max_plate_nesting=None,
        jit_compile=False,
        jit_options=None,
        ignore_jit_warnings=False,
        target_accept_prob=0.8,
        init_strategy=init_to_uniform,
        *,
        min_stepsize=1e-10,
        max_stepsize=1e10,
    ):

Import

from pyro.infer.mcmc import HMC

I/O Contract

Constructor Inputs

Parameter	Type	Required	Description
`model`	callable	No*	A Pyro model (stochastic function). Either `model` or `potential_fn` must be provided.
`potential_fn`	callable	No*	A Python callable that computes the potential energy given unconstrained parameters. Alternative to `model`.
`step_size`	float	No	Leapfrog step size. Defaults to `1`. Adapted during warmup if `adapt_step_size=True`.
`trajectory_length`	float	No	Total trajectory length. Number of leapfrog steps is computed as `trajectory_length / step_size`. Mutually exclusive with `num_steps`. Defaults to `2 * pi` if neither is specified.
`num_steps`	int	No	Fixed number of leapfrog integration steps per iteration. Mutually exclusive with `trajectory_length`.
`adapt_step_size`	bool	No	Whether to adapt step size during warmup. Defaults to `True`.
`adapt_mass_matrix`	bool	No	Whether to adapt the mass matrix during warmup. Defaults to `True`.
`full_mass`	bool	No	If `True`, uses a dense mass matrix. If `False`, uses a diagonal mass matrix. Defaults to `False`.
`transforms`	dict	No	Dictionary mapping sample site names to transforms. If `None`, inferred automatically from the model's constraints.
`max_plate_nesting`	int	No	Maximum depth of nested `pyro.plate` contexts.
`jit_compile`	bool	No	Whether to JIT-compile the potential function. Defaults to `False`.
`jit_options`	dict	No	Options for `torch.jit.trace`.
`ignore_jit_warnings`	bool	No	Whether to suppress JIT warnings. Defaults to `False`.
`target_accept_prob`	float	No	Target Metropolis acceptance probability for step size adaptation. Defaults to `0.8`.
`init_strategy`	callable	No	Strategy for initializing parameter values. Defaults to `init_to_uniform`.
`min_stepsize`	float	No	Minimum allowed step size during adaptation. Defaults to `1e-10`. Keyword-only.
`max_stepsize`	float	No	Maximum allowed step size during adaptation. Defaults to `1e10`. Keyword-only.

Outputs

Output	Type	Description
HMC kernel instance	`HMC`	A kernel object to be passed to the `MCMC` class. Provides `sample()`, `setup()`, and `logging()` methods.

Usage Examples

Basic HMC Sampling

import torch
import pyro
import pyro.distributions as dist
from pyro.infer.mcmc import HMC, MCMC

def model(data):
    mu = pyro.sample("mu", dist.Normal(0, 10))
    sigma = pyro.sample("sigma", dist.HalfNormal(10))
    with pyro.plate("data", len(data)):
        pyro.sample("obs", dist.Normal(mu, sigma), obs=data)

data = torch.randn(100) * 2 + 5

hmc_kernel = HMC(model, step_size=0.05, num_steps=20)
mcmc = MCMC(hmc_kernel, num_samples=1000, warmup_steps=500)
mcmc.run(data)

samples = mcmc.get_samples()
print(samples["mu"].mean())

HMC with Trajectory Length and Full Mass Matrix

from pyro.infer.mcmc import HMC, MCMC

hmc_kernel = HMC(
    model,
    trajectory_length=1.0,
    adapt_step_size=True,
    adapt_mass_matrix=True,
    full_mass=True,
    target_accept_prob=0.8,
    jit_compile=True,
)

mcmc = MCMC(hmc_kernel, num_samples=2000, warmup_steps=1000, num_chains=4)
mcmc.run(data)
mcmc.summary()

HMC with Custom Potential Function

import torch
from pyro.infer.mcmc import HMC, MCMC

def potential_fn(params):
    """Negative log probability of a 2D multivariate normal."""
    theta = params["theta"]
    return 0.5 * torch.sum(theta ** 2)

hmc_kernel = HMC(potential_fn=potential_fn, step_size=0.1, num_steps=10)
mcmc = MCMC(hmc_kernel, num_samples=500, warmup_steps=200)
mcmc.run()

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment