Samplers¤

Jim supports several JAX sampler backends behind a unified interface. You select one by passing a typed config object to Jim.

After jim.sample(), retrieve posterior samples with:

samples = jim.get_samples()  # dict[str, np.ndarray] keyed by parameter name

Sampler overview¤

flowMC¤

flowMC runs parallel MCMC chains enhanced by a normalizing flow that learns the posterior shape during training, then uses that learned geometry to make global proposals during production.

from jimgw.core.jim import Jim
from jimgw.samplers.config import FlowMCConfig

jim = Jim(
    likelihood,
    prior,
    sampler_config=FlowMCConfig(
        n_chains=1000,
        n_local_steps=100,
        n_global_steps=1000,
        n_training_loops=20,
        n_production_loops=10,
    ),
)

jim.sample()
samples = jim.get_samples()

Key parameters:

• n_chains — number of parallel MCMC chains.
• n_training_loops / n_production_loops — how many rounds of training (flow updates) and production (sample collection) to run.
• n_local_steps / n_global_steps — local MCMC steps and flow-proposal steps per loop.
• local_kernel — MCMC kernel for local steps; one of "MALA" (default), "HMC", or "GRW".
• parallel_tempering — parallel tempering settings; disabled by default. Enable with parallel_tempering=True (uses defaults), a plain dict of settings such as {"n_temperatures": 8}, or a ParallelTemperingConfig instance.

Repository: GW-JAX-Team/flowMC

References: Wong, K. W. K., Gabrié, M., Foreman-Mackey, D., "flowMC: Normalizing flow enhanced sampling package for probabilistic inference in JAX", arXiv:2211.06397, JOSS 8 (83) 5021 (2023). Wong, K. W. K., Isi, M., Edwards, T. D. P., "Fast Gravitational-wave Parameter Estimation without Compromises", arXiv:2302.05333, ApJ 958 129 (2023).

BlackJAX SMC¤

Sequential Monte Carlo (SMC) maintains a population of particles and gradually shifts them from the prior toward the posterior through a sequence of intermediate temperature steps.

Normalised-prior requirement — SMC computes a Bayesian evidence estimate and therefore requires a normalised prior. All built-in Jim priors are normalised. If you add custom constraints, check whether the resulting distribution is still normalised; if so, override is_normalized to return True. Jim raises a ValueError at construction if this condition is not met.

from jimgw.samplers.config import BlackJAXSMCConfig

jim = Jim(
    likelihood,
    prior,
    sampler_config=BlackJAXSMCConfig(
        n_particles=2000,
        n_mcmc_steps_per_dim=100,
    ),
)
jim.sample()
samples = jim.get_samples()

Key parameters:

• n_particles — particle ensemble size.
• n_mcmc_steps_per_dim — MCMC steps per dimension at each temperature step.
• target_ess_fraction — target ESS as a fraction of n_particles (default 0.9). The algorithm advances the temperature when the fraction of effectively contributing particles hits this threshold. Values in (0, 1] are valid when persistent_sampling=False; persistent sampling may exceed 1.0 because particles are recycled across steps. Only used with adaptive temperature selection (no effect with a fixed temperature_ladder).
• target_ess — target ESS as an absolute particle count. target_ess_fraction and target_ess are mutually exclusive; set one or the other, not both. When persistent_sampling=False, must be <= n_particles.
• persistent_sampling — whether to retain particles from all temperature steps (default True).
• temperature_ladder — explicit temperature schedule. If given, the sampler advances through this fixed ladder and ignores target_ess_fraction and target_ess.

Repository: blackjax-devs/blackjax

BlackJAX NS AW¤

Nested sampling with a bilby/dynesty-style adaptive differential-evolution acceptance-walk inner kernel.

Unit-cube requirement — this sampler works in the unit hypercube [0, 1]^n_dims. All parameters must be mapped into [0, 1] via sample transforms, which the CLI constructs automatically.

from jimgw.samplers.config import BlackJAXNSAWConfig

jim = Jim(
    likelihood,
    prior,
    sampler_config=BlackJAXNSAWConfig(
        n_live=1000,
        n_delete_frac=0.5,
        n_target=60,
        max_mcmc=5000,
        max_proposals=1000,
        termination_dlogz=0.1,
    ),
    sample_transforms=sample_transforms,
    likelihood_transforms=likelihood_transforms,
)
jim.sample()
samples = jim.get_samples()

Key parameters:

• n_live — number of live points; more live points → more accurate but slower.
• n_delete_frac — fraction of live points replaced per iteration (e.g. 0.5 replaces half the live points each step).
• n_target — target number of accepted proposals per walk.
• max_mcmc — maximum number of proposals before giving up on a dead point.

Reference: Prathaban, M., Yallup, D., Alvey, J., Yang, M., Templeton, W., Handley, W., "Gravitational-wave inference at GPU speed: A bilby-like nested sampling kernel within blackjax-ns", arXiv:2509.04336 (Sep 2025).

BlackJAX NSS¤

Nested sampling with a slice-sampling inner kernel. Unlike NS AW, it does not require a unit-cube prior and works in any bounded sampling space.

Normalised-prior requirement — NSS computes a Bayesian evidence estimate and therefore requires a normalised prior. All built-in Jim priors are normalised. If you add custom constraints, check whether the resulting distribution is still normalised; if so, override is_normalized to return True. Jim raises a ValueError at construction if this condition is not met.

from jimgw.samplers.config import BlackJAXNSSConfig

jim = Jim(
    likelihood,
    prior,
    sampler_config=BlackJAXNSSConfig(
        n_live=1000,
        n_delete_frac=0.5,
        num_inner_steps_per_dim=20,
        termination_dlogz=0.1,
        n_devices=1,
    ),
    sample_transforms=sample_transforms,
    likelihood_transforms=likelihood_transforms,
)
jim.sample()
samples = jim.get_samples()

Key parameters:

• n_live — number of live points.
• n_delete_frac — fraction of live points replaced per iteration.
• num_inner_steps_per_dim — slice-sampler steps per dimension per dead point; increase for strongly correlated posteriors.
• n_devices — number of local devices used to shard the live points and replacement chains. Both n_live and int(n_live * n_delete_frac) must be divisible by this value. The default 1 uses the unsharded kernel.

Repository: blackjax-devs/blackjax

References: Yallup, D., Prathaban, M., Alvey, J., Handley, W., "Parallel Nested Slice Sampling for Gravitational Wave Parameter Estimation", arXiv:2509.24949 (Sep 2025). Yallup, D., Kroupa, N., Handley, W., "Nested Slice Sampling", OpenReview (2025).

BlackJAX SwiG¤

Nested Slice within Gibbs (SwiG) applies covariance-shaped slice updates to named parameter blocks. With TransientLikelihoodFD, HeterodynedTransientLikelihoodFD, or MultibandedTransientLikelihoodFD, the likelihood provides a cache of generated waveform polarizations to the sampler. Blocks that affect waveform inputs refresh that cache, while detector-projection-only blocks reuse it. When the waveform model supports analytic distance scaling, luminosity distance is also factored out, so a d_L block reuses the cache; otherwise a d_L block refreshes it like any other waveform dependency. The likelihood determines the remaining dependencies after sample transforms, likelihood transforms, fixed parameters, and analytic marginalisation have been applied.

from jimgw.samplers.config import BlackJAXSwiGConfig

jim = Jim(
    likelihood,
    prior,
    sampler_config=BlackJAXSwiGConfig(
        blocks=[
            ["M_c", "q", "lambda_1", "lambda_2"],
            ["s1_z", "s2_z"],
            ["iota"],
            ["ra", "dec"],
            ["psi"],
            ["t_c"],
        ],
        n_live=512,
        n_delete_frac=0.125,
        num_inner_steps_per_dim=1,
        num_gibbs_sweeps=2,
        n_devices=1,
    ),
    likelihood_transforms=likelihood_transforms,
)

The cache-capable likelihood requires blocks to form an exact partition of the sampling parameters. A block that contains any waveform dependency refreshes the cache, so mixed blocks are safe but may be less efficient. Parameters that are analytically marginalised must not appear in the blocks — for example, the ["t_c"] block above assumes time_marginalization is disabled; see examples/GW170817_SwiG.py for a run with time and phase marginalisation enabled, where the t_c and phase_c blocks are dropped entirely.

Key parameters:

• blocks — ordered sampling-space parameter blocks for each Gibbs sweep.
• n_live / n_delete_frac — live-set size and replacement fraction, matching NSS.
• num_inner_steps_per_dim — random-direction slice steps per block dimension.
• num_gibbs_sweeps — complete block sweeps per particle replacement.
• n_devices — number of local devices used to shard the live points and replacement chains. The waveform cache remains local to each replacement chain; only live-point state participates in collectives.

Checkpointing and resuming¤

All samplers support checkpoint/resume so long-running jobs can survive interruptions. Set checkpoint_dir to a directory and checkpoint_interval to the minimum wall-clock seconds between writes:

from jimgw.samplers.config import BlackJAXSMCConfig

jim = Jim(
    likelihood,
    prior,
    sampler_config=BlackJAXSMCConfig(
        n_particles=2000,
        checkpoint_dir="./my_run",
        checkpoint_interval=600,  # write at most every 10 minutes
    ),
)
jim.sample()

The checkpoint is written atomically (checkpoint.pkl.tmp → checkpoint.pkl) so a mid-write crash never leaves a corrupt file. To resume after an interruption, construct the same config pointing at the same checkpoint_dir and call jim.sample() again — the sampler detects the existing file and picks up from the last saved state:

# resume — identical config, same checkpoint_dir
jim = Jim(
    likelihood,
    prior,
    sampler_config=BlackJAXSMCConfig(
        n_particles=2000,
        checkpoint_dir="./my_run",
        checkpoint_interval=600,
    ),
)
jim.sample()  # resumes from ./my_run/checkpoint.pkl

The same fields work identically for FlowMCConfig, BlackJAXNSAWConfig, BlackJAXNSSConfig and BlackJAXSwiGConfig.

Validation — setting checkpoint_interval > 0 without checkpoint_dir raises a ValidationError at config construction time.

Cross-sampler checkpoints — every checkpoint records which sampler backend wrote it. Pointing a different backend at that checkpoint_dir (or, for SMC, the same backend in a different mode — persistent/tempered/adaptive) is handled differently depending on the backend you point there:

• flowMC raises a ValueError immediately, before touching any sampler state.
• BlackJAX NS AW / NSS / SwiG / SMC log a warning ("incompatible or corrupt checkpoint … — starting fresh") and silently start a new run instead, the same way they handle a truly corrupt file.

If you switch sampler backends or SMC modes between runs, point checkpoint_dir at a new, empty directory to avoid silently discarding an existing checkpoint.

When using the CLI, checkpointing is enabled automatically (600 s, writing to output.dir). Set checkpoint_interval = 0 in the [sampler] block to opt out.

Periodic parameters¤

All samplers accept a periodic field to handle parameters that wrap around (e.g. angles). Pass a dict of parameter_name: (lower, upper) bounds:

config = FlowMCConfig(
    ...,
    periodic={"phase_c": (0.0, 6.283185), "psi": (0.0, 3.141593)},
)

BlackJAX NS AW operates in [0, 1] per dimension, so its periodic field takes a plain list of parameter names:

config = BlackJAXNSAWConfig(
    ...,
    periodic=["phase_c", "psi"],
)

Posterior samples¤

jim.get_samples() is the primary way to retrieve posterior samples.

samples = jim.get_samples()
# keys: prior parameter names + "log_likelihood"
samples["M_c"]             # np.ndarray — chirp mass in prior space
samples["log_likelihood"]  # np.ndarray — per-sample log-likelihood

Each backend's get_samples() returns equally-weighted posterior samples:

• NS AW / NSS: uses anesthetic's posterior_points() to resample the dead-point collection to equal-weight samples.
• SMC (persistent): resamples all-temperature particles weighted by the persistent-sampling weights.
• SMC (non-persistent): returns all final-temperature particles.
• flowMC: returns all production samples across all chains.

Pass n_samples to jim.get_samples() to further downsample uniformly without replacement:

samples = jim.get_samples(n_samples=2000)

Run diagnostics¤

jim.get_diagnostics() is a thin wrapper around the sampler's own get_diagnostics(), which returns a plain dict[str, Any].

diag = jim.get_diagnostics()

diag["n_likelihood_evaluations"]  # int   — total number of likelihood calls
diag["sampling_time"]             # float — wall-clock sampling time in seconds

Backend-specific keys:

# flowMC
diag["n_training_loops_actual"]         # int        — training loops run (may be less than configured)
diag["training_loss_history"]           # np.ndarray — normalizing-flow loss per epoch
diag["acceptance_training_local"]       # np.ndarray — local acceptance rate per training loop
diag["acceptance_training_global"]      # np.ndarray — global acceptance rate per training loop
diag["acceptance_production_local"]     # np.ndarray — local acceptance rate per production loop
diag["acceptance_production_global"]    # np.ndarray — global acceptance rate per production loop

# NS AW and NSS — also include evidence estimate
diag["n_iterations"]              # int   — number of nested-sampling steps
diag["log_Z"]                     # float — log Bayesian evidence
diag["log_Z_error"]               # float — standard deviation from 100 bootstrap samples

# SMC
diag["acceptance_history"]        # np.ndarray — mean acceptance rate at each step
diag["ess_history"]               # np.ndarray — ESS at each step
# Adaptive temperature only:
diag["n_iterations"]              # int        — number of temperature steps
diag["tempering_schedule"]        # np.ndarray — inverse temperature at each step
# Persistent sampling only:
diag["persistent_log_Z"]          # np.ndarray — cumulative log Z after each step
diag["log_Z"]                     # float      — final log Bayesian evidence

Writing your own sampler¤

This section is for advanced users who want to integrate a custom sampling backend with Jim. It requires familiarity with JAX and the Jim sampler internals.

Subclass Sampler, implement three methods and the sampler_name property, and register it:

• sampler_name — a human-readable backend name used in status and checkpoint logs.
• _sample(rng_key, initial_position) — run the sampler and store results. The base class wraps this in sample(), which also records sampling_time.
• get_samples() — return a dict with "samples" and "log_likelihood" keys.
• _get_diagnostics() — return a plain dict with diagnostic information. The base class wraps this in get_diagnostics(), which injects sampling_time.

from typing import Any, Literal, Optional
import numpy as np
from jimgw.samplers import register_sampler
from jimgw.samplers.base import Sampler
from jimgw.samplers.config import BaseSamplerConfig


class MyConfig(BaseSamplerConfig):
    type: Literal["my-sampler"] = "my-sampler"
    n_steps: int = 1000


class MySampler(Sampler):
    _config: MyConfig

    @property
    def sampler_name(self) -> str:
        return "My sampler"

    def __init__(self, *, n_dims, log_prior_fn, log_likelihood_fn,
                 log_posterior_fn, config: Optional[MyConfig] = None,
                 parameter_names=(), periodic=None):
        if config is None:
            config = MyConfig()
        super().__init__(n_dims=n_dims, log_prior_fn=log_prior_fn,
                         log_likelihood_fn=log_likelihood_fn,
                         log_posterior_fn=log_posterior_fn, config=config)
        self._result = None

    def _sample(self, rng_key, initial_position) -> None:
        # initial_position: shape (n_chains, n_dims), drawn from the prior by Jim.
        # ... run your sampler for self._config.n_steps steps ...
        self._result = np.asarray(initial_position)

    def get_samples(self) -> dict[str, np.ndarray]:
        if self._result is None:
            raise RuntimeError("call sample() first")
        return {
            "samples": self._result,
            "log_likelihood": np.zeros(self._result.shape[0]),
        }

    def _get_diagnostics(self) -> dict[str, Any]:
        if self._result is None:
            raise RuntimeError("call sample() first")
        return {
            "n_likelihood_evaluations": self._config.n_steps,
        }


register_sampler("my-sampler", lambda: MySampler)

Then pass MyConfig(n_steps=500) as sampler_config to Jim.