blackjax.mcmc.dynamic_hmc

Public API for the Dynamic HMC Kernel

Classes

DynamicHMCState

State of the dynamic HMC algorithm.

Functions

init(position, logdensity_fn, random_generator_arg)
build_kernel([integrator, divergence_threshold, ...])	Build a Dynamic HMC kernel where the number of integration steps is chosen randomly.
as_top_level_api([1], integration_steps_fn, 1, 10), ...)	Implements the (basic) user interface for the dynamic HMC kernel.
halton_sequence(→ float)	Generate the (i+1)-th element of the Halton sequence.

Module Contents

class DynamicHMCState

State of the dynamic HMC algorithm.

Adds a utility array for generating a pseudo or quasi-random sequence of number of integration steps.

position: blackjax.types.ArrayTree

logdensity: float

logdensity_grad: blackjax.types.ArrayTree

random_generator_arg: blackjax.types.Array

init(position: blackjax.types.ArrayLikeTree, logdensity_fn: Callable, random_generator_arg: blackjax.types.Array)

build_kernel(integrator: Callable = integrators.velocity_verlet, divergence_threshold: float = 1000, next_random_arg_fn: Callable = lambda key: ..., integration_steps_fn: Callable = lambda key: ..., build_proposal: Callable = hmc_proposal)

Build a Dynamic HMC kernel where the number of integration steps is chosen randomly.

Parameters:

• integrator – The symplectic integrator to use to integrate the Hamiltonian dynamics.

• divergence_threshold – Value of the difference in energy above which we consider that the transition is divergent.

• next_random_arg_fn – Function that generates the next random_generator_arg from its previous value.

• integration_steps_fn – Callable with signature (random_generator_arg, *integration_steps_params) -> int that draws the number of integration steps for a single transition. Extra positional arguments beyond random_generator_arg are supplied at call time via integration_steps_params on the inner kernel, so tunable parameters (e.g. average number of steps, distribution bounds) can be adapted without rebuilding the kernel.

• build_proposal – A callable with signature (integrator, kinetic_energy, step_size, num_integration_steps, divergence_threshold) -> generate that builds the proposal function. Defaults to hmc_proposal() (standard endpoint HMC).

Returns:

• A kernel that takes a rng_key and a Pytree that contains the current state

• of the chain and that returns a new state of the chain along with

• information about the transition.

as_top_level_api(logdensity_fn: Callable, step_size: float, inverse_mass_matrix: blackjax.types.Array, *, divergence_threshold: int = 1000, integrator: Callable = integrators.velocity_verlet, next_random_arg_fn: Callable = lambda key: ..., integration_steps_fn: Callable = lambda key: ..., integration_steps_params: tuple = (), build_proposal: Callable = hmc_proposal) → blackjax.base.SamplingAlgorithm

Implements the (basic) user interface for the dynamic HMC kernel.

Parameters:

• logdensity_fn – The log-density function we wish to draw samples from.

• step_size – The value to use for the step size in the symplectic integrator.

• inverse_mass_matrix – The value to use for the inverse mass matrix when drawing a value for the momentum and computing the kinetic energy.

• divergence_threshold – The absolute value of the difference in energy between two states above which we say that the transition is divergent. The default value is commonly found in other libraries, and yet is arbitrary.

• integrator – (algorithm parameter) The symplectic integrator to use to integrate the trajectory.

• next_random_arg_fn – Function that generates the next random_generator_arg from its previous value.

• integration_steps_fn – Callable with signature (random_generator_arg, *integration_steps_params) -> int that draws the number of integration steps for a single transition.

• integration_steps_params – Extra positional arguments unpacked into integration_steps_fn after random_generator_arg on every step. Use this to pass tunable parameters (e.g. (avg_num_integration_steps,) or (lower_bound, upper_bound)) without rebuilding the kernel. Defaults to () so that a plain 1-arg integration_steps_fn works unchanged.

• build_proposal – A callable with signature (integrator, kinetic_energy, step_size, num_integration_steps, divergence_threshold) -> generate that builds the proposal function. Defaults to hmc_proposal() (standard endpoint HMC). Pass multinomial_hmc_proposal() for multinomial trajectory sampling.

Return type:

A SamplingAlgorithm.

halton_sequence(i: blackjax.types.Array, max_bits: int = 10) → float

Generate the (i+1)-th element of the Halton sequence.

Warning: max_bits should be less than the bit width of i.dtype to prevent integer overflow (e.g., max_bits <= 63 for int64).