BlackJAX

Blackjax provides a wide range of sampling and inference algorithms categorized into several families:

MCMC (Markov Chain Monte Carlo)

Includes Hamiltonian Monte Carlo (hmc), No-U-Turn Sampler (nuts), Metropolis-Adjusted Langevin Algorithm (mala), and various Random-walk Metropolis-Hastings (rmh) implementations.

MCLMC Family

Includes Microcanonical Langevin Monte Carlo (mclmc) and its adjusted variants (adjusted_mclmc).

Laplace-preconditioned Family

HMC and its variants (Dynamic, Multinomial) that use Laplace approximation preconditioning (laplace_hmc).

Stochastic Gradient MCMC

Algorithms for large datasets like Stochastic Gradient Langevin Dynamics (sgld) and Stochastic Gradient HMC (sghmc).

Sequential Monte Carlo (SMC)

Includes Tempered SMC (tempered_smc) and various persistent-particle SMC implementations.

Variational Inference (VI)

Includes Mean-field (meanfield_vi), Full-rank (fullrank_vi), and Pathfinder (pathfinder) methods.

Adaptation / Warmup

Tools to tune step-sizes and mass matrices, such as window_adaptation and chees_adaptation.

Diagnostics & Utilities

• blackjax.ess: Effective Sample Size
• blackjax.rhat: Potential Scale Reduction (R̂)
• run_inference_algorithm: A lax.scan-based inference loop utility for speed.

Funsor is a library of functional tensors that allows for exact marginalisation of discrete latent variables. This is a powerful complement to BlackJAX because it enables gradient-based samplers like NUTS to work on models containing discrete variables (e.g., Gaussian Mixture Models, Hidden Markov Models) by turning sums over discrete states into differentiable JAX operations.

To use this pattern, you define your log-density using Funsor primitives, which allows the discrete variables to be marginalised out before passing the resulting continuous log-density to a BlackJAX kernel.

import blackjax
import funsor
import funsor.ops as ops
from funsor.domains import Bint
from funsor.tensor import Tensor
from funsor.terms import Variable
from funsor.jax.distributions import Categorical, Normal

# 1. Set backend
funsor.set_backend("jax")

# 2. Define log-density using Funsor primitives
def gmm_logdensity(position):
    # ... (use Tensor, Variable, and distributions to define log_p) ...
    # Example of marginalisation:
    # log_marginal_n = log_joint.reduce(ops.logaddexp, "z")
    return log_p.data

# 3. Use BlackJAX for inference
adapt = blackjax.window_adaptation(blackjax.nuts, gmm_logdensity)
(last_state, parameters), _ = adapt.run(rng_key, position0, num_steps=1000)
kernel = blackjax.nuts(gmm_logdensity, **parameters).step

When using blackjax.progress_bar(), be aware of the following technical behaviors:

• Checkpoint + Gradient: If using jax.checkpoint with differentiation, the callback fires twice per logical step (primal + recompute), making step counts appear roughly doubled.
• Multi-device sharding: Under jax.shard_map, the callback fires once per device per step. This increases host dispatch overhead, and the bar may reach 100% while slower shards are still processing.
• functools.partial bypass: Using functools.partial(jax.lax.scan, ...) captured before entering the progress bar context will silently bypass the bar without error.
• Shared output_file: If multiple contexts use the same output_file path, writes will interleave and corrupt the file. Use unique paths for different contexts.
• Nesting with jaxtap.record(): If a jaxtap.record() context is opened inside a progress_bar() context, the inner context takes precedence. The progress bar will be silent for scans run within that inner block.

Avoid using Python for loops to iterate over MCMC steps inside a JIT-compiled function. A Python loop causes JAX to unroll the entire computation graph at trace time, leading to $O(n)$ compilation costs and massive memory usage for large step counts.

Instead, use jax.lax.scan. This lowers to a single XLA WhileOp, keeping compilation cost constant ($O(1)$) regardless of the number of steps. For best results, wrap the function containing the lax.scan in an outer @jax.jit to cache the full trace.

# ✗ Python loop inside jit — unrolls at trace time, O(n) compile cost
@jax.jit
def run_python_loop(rng_key, initial_state, n_steps):
    state = initial_state
    for key in jax.random.split(rng_key, n_steps):
        state, _ = kernel(key, state)
    return state

# ✓ lax.scan inside jit — O(1) compile cost, compact XLA WhileOp
@jax.jit
def run_scan(rng_key, initial_state, n_steps):
    def step(state, key):
        state, info = kernel(key, state)
        return state, info
    return jax.lax.scan(step, initial_state, jax.random.split(rng_key, n_steps))

Metropolis-Within-Gibbs (MWG) sampling is used when you want to sample from a joint distribution $p(x, y)$ by alternately updating components (e.g., $x$ and $y$) using different MCMC kernels.

When implementing MWG in BlackJAX, a critical requirement is ensuring that the AlgorithmState of each component reflects the most recent positions of all other components. Because BlackJAX kernels are pure, updating one component (e.g., $x$) changes the conditional log-density for the next component (e.g., $y$).

To maintain correctness, you must manually update the log_probability in the AlgorithmState of the next component before calling its step function. This is achieved by using the blackjax.algorithm.init() function, passing the current position and the new conditional logdensity_fn.

# Example logic for updating state between component steps
# 1. Update x using its conditional log-density
state["x"] = mwg_init_x(
    position=state["x"].position,
    logdensity_fn=logdensity_x
)
state["x"], _ = mwg_step_fn_x(
    rng_key=rng_key_x,
    state=state["x"],
    logdensity_fn=logdensity_x,
    **parameters["x"]
)

# 2. Update y using the NEW position of x
# We must re-initialize state["y"] so its log_probability matches p(y | x_new)
def logdensity_y(y): return logdensity(y=y, x=state["x"].position)

state["y"] = mwg_init_y(
    position=state["y"].position,
    logdensity_fn=logdensity_y
)
state["y"], _ = mwg_step_fn_y(
    rng_key=rng_key_y,
    state=state["y"],
    logdensity_fn=logdensity_y,
    **parameters["y"]
)

Every algorithm in blackjax/mcmc/, blackjax/vi/, and blackjax/sgmcmc/ follows a consistent three-layer structure. This allows you to choose between manual state management or a high-level convenience wrapper.

1. init: Creates the initial algorithm state.
   • Signature: (position, logdensity_fn, *, rng_key=None, **kwargs) -> State
   • Note: If initialization requires randomness (e.g., sampling momentum), pass rng_key as a keyword argument.
2. build_kernel: Returns a specialized kernel function via a closure. This is the idiomatic way to configure an algorithm (e.g., setting step size or integrator type) once to avoid recompilation during sampling.
   • Signature: (**params) -> kernel_fn
   • Kernel Signature: kernel(rng_key, state, logdensity_fn, **params) -> (new_state, info)
3. as_top_level_api: A convenience wrapper that binds the logdensity_fn and parameters once, returning a SamplingAlgorithm object.
   • Signature: (logdensity_fn, **params) -> SamplingAlgorithm

SamplingAlgorithm is a NamedTuple containing (init, step), where step(rng_key, state) -> (new_state, info).

# Example conceptual usage of the three layers

# 1. Using init and build_kernel manually
state = hmc.init(initial_position, logdensity_fn, rng_key=key)
kernel_fn = hmc.build_kernel(step_size=0.1)
new_state, info = kernel_fn(key, state, logdensity_fn)

# 2. Using the top-level API
sampling_alg = hmc.as_top_level_api(logdensity_fn, step_size=0.1)
state = sampling_alg.init(initial_position)
new_state, info = sampling_alg.step(key, state)

BlackJAX follows a Markovian approach, where the current state contains all the information required to compute the next iteration. The interface differs depending on the algorithm type:

Sampling Algorithms (MCMC, SMC, SGMCMC)

These consist of an initializer and an iterator. The user initializes a state and then calls step to transition to a new state.

Approximate Inference Algorithms (VI)

These consist of an initializer, an iterator, and a sampler. In addition to step, they provide a sample method to draw samples from the learned distribution.

# Sampling Algorithm Pattern
sampling_algorithm = blackjax.nuts(logdensity_fn, step_size, inverse_mass_matrix)
state = sampling_algorithm.init(initial_position)
new_state, info = sampling_algorithm.step(rng_key, state)

# Approximate Inference Algorithm Pattern
approx_inf_algorithm = blackjax.pathfinder(logdensity_fn)
state = approx_inf_algorithm.init(initial_position)
new_state, info = approx_inf_algorithm.step(rng_key, state)
position_samples = approx_inf_algorithm.sample(rng_key, state, num_samples)