blackjax.mcmc.composed.trajectory_length

GIST instance (b): self-tuning trajectory length (no-U-turn, GIST paper section 5) – NOT NUTS’s recursive doubling.

The tuning parameter is alpha = L, the number of leapfrog steps at a fixed step size h. L is drawn uniformly from [Lo(theta, rho) : U(theta, rho)] where U(theta, rho) is the number of leapfrog steps until the no-U-turn condition first fires (a single forward while_loop rollout, materially simpler than NUTS’s recursive-doubling tree, section 2.2.2) and Lo shifts the lower end of that range by a fixed path fraction psi. Because g = identity on L, the reversibility of the resulting kernel is automatic (Corollary 4); the only thing the acceptance ratio has to account for is that the forward and reverse draws are uniform over different-width intervals containing the same L (section 2.2.4).

Forward-rollout caching

The forward U(theta, rho) search already integrates every state on [1 : U] one leapfrog step at a time (the while_loop in num_steps_to_uturn()); the selected proposal is always one of those states (some L <= U). Rather than throw them away and re-integrate L steps from scratch via trajectory.static_integration (the original implementation), _tuning_parameter_fn() buffers every rolled-out state as it goes (_num_steps_to_uturn_with_rollout()) and threads the buffer to _apply_fn() as GIST’s aux, so building the proposal is an O(1) gather rather than an O(L) re-integration. Per transition this drops the cost from roughly U + L + U_rev leapfrog/gradient evaluations (~2.75 U empirically, averaging over L’s distribution) to roughly U + U_rev (~1.25 U). The reverse rollout U_rev(theta', rho') (section 2.2.4) has nothing to gather from (its own states are never selected as the proposal) and is left as an ordinary, unbuffered num_steps_to_uturn() call.

The buffer is max_num_steps copies of the full leapfrog state (position, momentum, logdensity, logdensity_grad) – the same asymptotic memory scale as the states NUTS’s own trajectory-doubling machinery holds live during tree building (blackjax.mcmc.trajectory), and it is allocated fresh on every kernel call (not carried across transitions), so steady-state memory is O(max_num_steps * dim) on top of the chain state itself, not O(num_samples * max_num_steps * dim). This makes max_num_steps (default 1024, mirroring NUTS’s max_num_doublings=10 cap of 1023 steps) a memory knob as well as a compute-cap knob: sizing it much larger than the rollout lengths a target actually needs pays for unused buffer capacity on every transition, so prefer capping it near the largest U the warmup/pilot run actually observes rather than leaving the default untouched on very high-dimensional targets.

References

[1]

Bou-Rabee, Carpenter, Marsden, “GIST: Gibbs self-tuning for locally adaptive Hamiltonian Monte Carlo”, arXiv:2404.15253, section 5 (p.20-25), Algorithm 2 (p.21), eq. 33 (the no-U-turn condition), eq. 34-35 (the step distributions).

Attributes

init

Classes

GISTTrajectoryLengthInfo

Additional information for a gist_trajectory_length transition.

Functions

num_steps_to_uturn(→ collections.abc.Callable)	Build the U(theta, rho) forward-rollout function, section 2.2.2.
build_kernel(→ collections.abc.Callable)	Build a gist_trajectory_length kernel.
as_top_level_api(→ blackjax.base.SamplingAlgorithm)	blackjax.gist_trajectory_length -- GIST self-tuning path length

Module Contents

init

class GISTTrajectoryLengthInfo

Additional information for a gist_trajectory_length transition.

momentum, tuning_parameter, is_accepted, is_divergent, acceptance_rate, energy, num_integration_steps

Same convention as GISTStepSizeInfo (flat extension, not nesting). tuning_parameter is L itself (an integer, not a compound PyTree, since g = identity is the only tuning-parameter component here).

num_steps_to_uturn_forward, num_steps_to_uturn_reverse

U(theta, rho) = M and U(theta', rho') = N, section 2.2.2 – the (possibly capped) leapfrog-step counts to the no-return condition, forward from the current state and from the proposal.

is_no_return_rejected

True when L fell outside [Lo(theta', rho'), N] – the paper’s own “no-return” rejection category (Fig. 5, section 2.2.4), tracked separately from an ordinary energy-based Metropolis rejection.

momentum: blackjax.types.ArrayTree

tuning_parameter: blackjax.types.Array

is_accepted: blackjax.types.Array

is_divergent: blackjax.types.Array

acceptance_rate: blackjax.types.Array

energy: float

num_integration_steps: blackjax.types.Array

num_steps_to_uturn_forward: blackjax.types.Array

num_steps_to_uturn_reverse: blackjax.types.Array

is_no_return_rejected: blackjax.types.Array

num_steps_to_uturn(integrator: collections.abc.Callable, step_size: float, metric: blackjax.mcmc.metrics.Metric, max_num_steps: int) → collections.abc.Callable

Build the U(theta, rho) forward-rollout function, section 2.2.2.

A single forward while_loop, one leapfrog step at a time, checking the sign of the (metric-corrected) no-return condition after every step – no tree/doubling data structure, no recursion, no sub-U-turn bookkeeping, unlike NUTS’s 1sub-U-turn criterion (section 3.5).

The dot product pairs the position displacement with the raw momentum rho (GIST paper eq. 33), so the criterion generalizes correctly to a non-identity inverse_mass_matrix (diagonal / dense / low-rank) in an affine-equivariant manner: with p ~ N(0, G^-1) and K(p) = ½ pᵀGp, whitening φ = G^{-1/2}θ and p_φ = G^{1/2}p gives (Δθ)ᵀp = Δφᵀp_φ, which is exactly invariant under any change of basis G. This is a different criterion from NUTS’s momentum-sum check_turning, which uses the momentum sum (not position displacement) in the pairing. This changes no formula for the identity-metric case the paper’s own experiments use.

max_num_steps is a hard cap, exactly analogous to NUTS’s max_num_doublings (except here it bounds a linear rollout, not 2**max leapfrog steps). Capping is not an approximation: as long as p(L|theta,rho) and the acceptance ratio consistently use this capped U, the resulting p(L|theta,rho) is still an exact, well-defined, strictly-positive conditional density (the GIST reversibility guarantee, Theorem 3, does not care that U is a capped version of the “true” U-turn step count).

Parameters:

• integrator – Symplectic integrator used for the one-leapfrog-at-a-time rollout.

• step_size – h, the fixed step size (not GIST-adapted in this instance).

• metric – The (already-resolved) Metric.

• max_num_steps – Hard cap on the rollout length.

Returns:

• uturn_fn(state: IntegratorState, logdensity_fn) -> Array, the

• (possibly capped) number of leapfrog steps to the no-return condition.

build_kernel(integrator: collections.abc.Callable = integrators.velocity_verlet, divergence_threshold: float = 1000, path_fraction: float = 0.5, max_num_steps: int = 1024) → collections.abc.Callable

Build a gist_trajectory_length kernel.

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.

• path_fraction – psi in [0, 1], section 2.2.3. Default 0.5 per the paper’s own recommendation (comparable leapfrog-step counts to NUTS; psi=0 is the simpler eq. 34 special case).

• max_num_steps – Hard cap on each U-turn rollout (forward and reverse); analogous to NUTS’s max_num_doublings, but bounds a linear rollout here, not 2**max leapfrog steps. Also sizes the forward-rollout buffer (module docstring, “Forward-rollout caching”) – O(max_num_steps) leapfrog states are allocated per transition.

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: collections.abc.Callable, inverse_mass_matrix: blackjax.mcmc.metrics.MetricTypes, step_size: float, *, path_fraction: float = 0.5, max_num_steps: int = 1024, divergence_threshold: float = 1000, integrator: collections.abc.Callable = integrators.velocity_verlet) → blackjax.base.SamplingAlgorithm

blackjax.gist_trajectory_length – GIST self-tuning path length (no-U-turn condition, section 2.2; NOT NUTS’s recursive doubling).

Examples

A new gist_trajectory_length kernel can be initialized and used with the following code:

gist_trajectory_length = blackjax.gist_trajectory_length(
    logdensity_fn, inverse_mass_matrix, step_size=0.1
)
state = gist_trajectory_length.init(position)
new_state, info = gist_trajectory_length.step(rng_key, state)

Parameters:

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

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

• step_size – h, fixed (not GIST-adapted in this instance; see the module docstring for composing with gist_step_size).

• path_fraction – psi in [0, 1], section 2.2.3. Default 0.5 per the paper’s own recommendation (comparable leapfrog-step counts to NUTS; psi=0 is the simpler eq. 34 special case).

• max_num_steps – Hard cap on each U-turn rollout (forward and reverse); analogous to NUTS’s max_num_doublings, but bounds a linear rollout here, not 2**max leapfrog steps – size accordingly (NUTS’s default of 10 doublings caps at 1023 steps; a directly-comparable cap here is max_num_steps ~= 1024). Also sizes the forward-rollout buffer (module docstring, “Forward-rollout caching”) – O(max_num_steps) leapfrog states are allocated per transition, so an oversized cap on a high-dimensional target trades unused memory for compute.

• divergence_threshold – The absolute value of the difference in energy between two states above which we say that the transition is divergent.

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

Return type:

A SamplingAlgorithm.