Skip to content

← JAX Stochasticity step 14 of 25

Hard Primitives

Leapfrog with a Constant Force

Jump to the task

Why this matters

Leapfrog updates momentum for half a step, position for a full step, then momentum for another half step. Hamiltonian Monte Carlo uses these updates to construct proposals that follow gradients of the log target density. A complete sampler also refreshes momentum and applies an acceptance rule; integration alone is not an HMC sampler.

This exercise practices the update order with unit mass and a supplied constant force grad_log_prob. It is a simplified dynamics calculation. A Gaussian target does not have a constant log-density gradient: for a standard normal target the gradient at position x is -x, which changes as the position changes.

Worked mini-example

Start at position zero with zero momentum. With constant force 1 and step size 0.5, one step is:

x, p = 0.0, 0.0
force, step = 1.0, 0.5
p_half = p + 0.5 * step * force  # 0.25
x_new = x + step * p_half       # 0.125
p_new = p_half + 0.5 * step * force  # 0.5

The second step starts from (0.125, 0.5) and ends at (0.5, 1.0). The position update must use the half-updated momentum, not the original momentum or the final momentum.

Problem

Implement leapfrog_step(x, p, grad_log_prob, step_size, n_steps). For each iteration, apply these updates in order:

p_half = p + 0.5 * step_size * grad_log_prob
x_new = x + step_size * p_half
p_new = p_half + 0.5 * step_size * grad_log_prob

Carry (x_new, p_new) into the next iteration. All inputs are scalars. Position, momentum, force and step size are finite floats. n_steps is a nonnegative integer-valued number representable as int32 and may arrive as a float. Assume the arithmetic remains finite in float32. Return [final_x, final_p] as a float32 array of shape (2,). Zero steps leave the state unchanged. A negative step size integrates backward for this constant-force system.

Use lax.fori_loop to practice carrying state through a JAX loop. For checking your result, constant force also gives a closed form: with t = n_steps * step_size, final momentum is p + t * force and final position is x + t * p + 0.5 * t**2 * force, up to rounding. The output fixtures cannot distinguish that formula from an iterative implementation.

Common pitfalls

  • Both half momentum updates belong in each full step. Omitting one changes the specified trajectory. It does not prove that every such alternative update fails to preserve volume.
  • The force here is the gradient of the log target, so momentum updates add it. Using the gradient of potential energy instead would reverse the sign, because potential energy is negative log density.
  • General leapfrog evaluates the force at the current position for the first half step and at the new position for the second. Reusing this supplied constant models a linear potential, not a harmonic oscillator. Such a constant-force example is not a normalizable target on the whole real line.
  • Cast the loop bound with jnp.int32(n_steps). A Python loop can work eagerly or with a concrete static bound, but a traced dynamic bound needs a JAX loop. fori_loop may lower to scan or while depending on whether its bounds are known during tracing; it is not guaranteed to become one machine instruction.