bijx.IMH

class bijx.IMH

Bases: Module

Independent Metropolis-Hastings sampler.

Implements the Independent Metropolis-Hastings algorithm for sampling from a target distribution using an independent proposal distribution. This is particularly useful when the proposal distribution is a good approximation to the target, such as using a normalizing flow as the proposal.

The algorithm generates proposals independently from the current state, then accepts or rejects them based on the Metropolis criterion:

\[\alpha = \min\left(1, \frac{p(x')q(x)}{p(x)q(x')}\right)\]

where \(p(x)\) is the target density and \(q(x)\) is the proposal density.

Key differences from blackjax.irmh:

• Sampler returns both samples and their log probabilities

• Integration with bijx distribution/bijection ecosystem

• Flax NNX module system compatibility

Example

>>> target_log_prob = lambda x: -0.5 * jnp.sum(x**2)  # Standard normal
>>> proposal = bijx.IndependentNormal(event_shape=(2,))
>>> sampler = bijx.IMH(proposal, target_log_prob)
>>> initial_state = sampler.init(key)
>>> new_state, info = sampler.step(key, initial_state)

__init__(sampler, target_log_prob)

Initialize Independent Metropolis-Hastings sampler.

Parameters:

• sampler – Proposal distribution implementing sample() method that returns (position, log_prob) tuples.

• target_log_prob (Callable) – Function computing log probability density of target distribution.

Methods

init(rng)	Initialize the sampler state.
propose(rng)	Generate a proposal state.
step(rng, state)	Perform one step of Independent Metropolis-Hastings.

propose(rng)

Generate a proposal state.

Samples from the proposal distribution and evaluates both proposal and target log probabilities at the sampled position.

Parameters:

rng – Random key for sampling.

Returns:

IMHState containing the proposal position and log probabilities.

init(rng)

Initialize the sampler state.

Parameters:

rng – Random key for initialization.

Returns:

Initial IMHState by proposing from the proposal distribution.

step(rng, state)

Perform one step of Independent Metropolis-Hastings.

Generates a proposal and applies the Metropolis acceptance criterion to decide whether to move to the new state or remain at the current one.

Parameters:

• rng – Random key for the step.

• state – Current IMHState of the chain.

Returns:

• new_state: Updated IMHState (either proposal or unchanged)

• info: IMHInfo with acceptance decision and diagnostics

Return type:

Tuple of (new_state, info) where

Note

The acceptance probability is computed as: \(\alpha = \min(1, \exp(\log p(x') - \log q(x') - \log p(x) + \log q(x)))\)