Skip to content

← JAX Stochasticity step 12 of 25

Medium Primitives

Metropolis-Hastings Step

Jump to the task

Why this matters

A Metropolis step can accept a move to a less likely state. This is how a Markov chain explores a distribution instead of only climbing toward its highest-density point. If a proposal is rejected, the current state appears again in the chain; rejection is not a missing sample.

This exercise handles the acceptance decision for a supplied proposal. Assume the proposal distribution is symmetric, such as a Gaussian random walk with the same covariance at every position. Then the forward and backward proposal densities cancel, leaving the target-density ratio. A Gaussian centered on the current state is not enough for symmetry if its covariance changes with position.

Worked mini-example

import jax
import jax.numpy as jnp

x, proposal_x = 0.0, 1.0
log_p_current, log_p_proposal = -1.0, -0.5
key = jax.random.PRNGKey(0)
log_alpha = jnp.minimum(0.0, log_p_proposal - log_p_current)  # 0
u = jax.random.uniform(key)  # Scalar in [0, 1).
accepted = jnp.log(u) < log_alpha
new_x = jnp.where(accepted, proposal_x, x)
result = jnp.array([new_x, accepted.astype(jnp.float32)])  # [1, 1]

The proposed target density is larger here, so acceptance has probability one. A log-density difference of -1 would instead give acceptance probability exp(-1), about 0.368.

Problem

Implement mh_step(seed, x, log_p_current, proposal_x, log_p_proposal):

  1. Create jax.random.PRNGKey(int(seed)).
  2. Set log_alpha = min(0, log_p_proposal - log_p_current).
  3. Draw one scalar with jax.random.uniform(key) and accept exactly when jnp.log(u) < log_alpha.
  4. Return [proposal_x, 1.0] if accepted, otherwise [x, 0.0], as a float32 array of shape (2,).

seed is a nonnegative integer-valued number supplied as a scalar. The states and target log densities are finite scalars; assume their arithmetic remains finite in float32. The fixtures use the specified random draw. Do not generate the proposal itself.

Common pitfalls

  • Equal target log densities also give acceptance probability one. With an asymmetric proposal, target densities alone are insufficient.
  • In exact arithmetic, u < exp(log_alpha) is an equivalent comparison. Since log_alpha <= 0, that exponential cannot overflow, but it can underflow for a very negative value. The requested log comparison avoids computing it.
  • JAX's uniform draw can be zero. Then log(u) is negative infinity, which is below any finite log acceptance threshold. Use the specified strict comparison to match the exercise's finite-precision convention.
  • Keep the state when rejecting. The returned flag records the decision, even if the proposed and current states happen to have equal values.
  • Mark seed static if wrapping this implementation in jax.jit, because the Python int(seed) conversion needs a concrete value.