Inference¶

Pure kernels, adaptation utilities and compiled runners. See Inference for how they fit together.

Runners¶

probjax.inference.MCMC ¶

Bases: WithProgressBarAPI, RunnerMixin

Compiled standard execution for an MCMC kernel.

The static methods are standalone compiled primitives. Instance methods apply the runner's kernel, collection, and progress configuration.

Source code in probjax/inference/mcmc_runner.py

class MCMC(WithProgressBarAPI, RunnerMixin):
    """Compiled standard execution for an MCMC kernel.

    The static methods are standalone compiled primitives. Instance methods
    apply the runner's kernel, collection, and progress configuration.
    """

    _default_tracked_stats = ("logdensity", "acceptance_rate")
    _print_rate = 50

    def __init__(
        self,
        kernel: MarkovKernel,
        verbose: bool = False,
        tracked_stats: Optional[Tuple[str, ...]] = None,
        collect_state: Tuple[str, ...] = (),
        collect_info: Tuple[str, ...] = (),
    ) -> None:
        self.kernel = kernel
        self.verbose = verbose
        self.tracked_stats = (
            self._default_tracked_stats if tracked_stats is None else tracked_stats
        )
        self.collect_state = collect_state
        self.collect_info = collect_info

    @staticmethod
    @partial(
        jax.jit,
        static_argnames=(
            "kernel",
            "num_steps",
            "collect_state",
            "collect_info",
        ),
    )
    def run_kernel(
        key: RngKey,
        kernel: MarkovKernel,
        state: State,
        num_steps: int,
        params: Params,
        args: Optional[Tuple] = None,
        *,
        collect_state: Tuple[str, ...] = (),
        collect_info: Tuple[str, ...] = (),
    ) -> MCMCResult:
        """Run a compiled kernel scan without constructing a runner."""
        keys = jax.random.split(key, num_steps)
        xs = keys if args is None else (keys, args)

        one_step = _mcmc_run_body(kernel, params, args, collect_state, collect_info)

        state, (states, info) = jax.lax.scan(one_step, state, xs)
        trace = states if collect_state else None
        diagnostics = info if collect_info else None
        return MCMCResult(state, trace, diagnostics)

    # Re-exposed from mcmc.base so kernels can sample without importing the runner.
    sample_kernel = staticmethod(sample_kernel)

    adaptor_warmup = staticmethod(adaptor_warmup)
    adapt_step = staticmethod(adapt_step)

    def _prepare_params(self, state: State, params: Optional[Params]) -> Params:
        return self.kernel.init_params(state) if params is None else params

    def run(
        self,
        key: RngKey,
        state: State,
        num_steps: int,
        params: Optional[Params] = None,
        args: Optional[Tuple] = None,
    ) -> MCMCResult:
        """Run ``num_steps`` transitions with this runner's configuration."""
        params = self._prepare_params(state, params)
        if not self.verbose:
            return self.run_kernel(
                key,
                self.kernel,
                state,
                num_steps,
                params,
                args,
                collect_state=self.collect_state,
                collect_info=self.collect_info,
            )

        keys = jax.random.split(key, num_steps)
        xs = keys if args is None else (keys, args)

        one_step = _mcmc_run_body(
            self.kernel,
            params,
            args,
            self.collect_state,
            self.collect_info,
            stats_fn=self._extract_stats,
        )

        state, (_, states, info) = self._verbose_scan(
            one_step,
            state,
            xs,
            num_steps,
            stats_fn=lambda _carry, y: y[0],
        )
        trace = states if self.collect_state else None
        diagnostics = info if self.collect_info else None
        return MCMCResult(state, trace, diagnostics)

    def sample(
        self,
        key: RngKey,
        state: State,
        num_samples: int,
        params: Optional[Params] = None,
        thin: int = 1,
        args: Optional[Tuple] = None,
    ) -> MCMCResult:
        """Collect positions after every ``thin`` transitions."""
        params = self._prepare_params(state, params)
        if not self.verbose:
            return self.sample_kernel(
                key,
                self.kernel,
                state,
                num_samples,
                params,
                args,
                thin=thin,
                collect_info=self.collect_info,
            )

        keys = jax.random.split(key, (num_samples, thin))
        step_args = _reshape_thin_args(args, num_samples, thin)
        xs = keys if step_args is None else (keys, step_args)

        collect_one = _mcmc_collect_body(
            self.kernel,
            params,
            step_args,
            self.collect_info,
            stats_fn=self._extract_stats,
        )

        state, (_, samples, info) = self._verbose_scan(
            collect_one,
            state,
            xs,
            num_samples,
            stats_fn=lambda _carry, y: y[0],
        )
        return MCMCResult(state, samples, info if self.collect_info else None)

    def adapt(
        self,
        key: RngKey,
        adaptor: Adaptor,
        state: State,
        params: Params,
        num_steps: int,
        args: Optional[Tuple] = None,
        *,
        collect: bool = False,
    ) -> AdaptationResult:
        """Fit kernel parameters with a composable adaptor."""
        return self.adaptor_warmup(
            key,
            self.kernel,
            adaptor,
            state,
            params,
            num_steps,
            args,
            collect=collect,
        )

    @staticmethod
    def _resolve_policy(policy: Union[Adaptor, Warmup]) -> Warmup:
        """Accept a bare adaptor (auto-wrapped) or a warmup policy."""
        if isinstance(policy, Adaptor):
            return as_warmup(policy)
        if isinstance(policy, Warmup):
            return policy
        raise TypeError(
            f"adaptation must be an Adaptor or a Warmup, got {type(policy).__name__}."
        )

    @staticmethod
    def warmup_kernel(
        key: RngKey,
        kernel: MarkovKernel,
        adaptation: Union[Adaptor, Warmup],
        state: State,
        params: Params,
        num_steps: int,
    ) -> WarmupResult:
        return MCMC._resolve_policy(adaptation).run(
            key, kernel, state, params, num_steps
        )

    def warmup(
        self,
        key: RngKey,
        adaptation: Union[Adaptor, Warmup, None] = None,
        state: Optional[State] = None,
        params: Optional[Params] = None,
        num_steps: Optional[int] = None,
        **adaptation_kwargs,
    ) -> WarmupResult:
        """Run a specialized state-producing warmup procedure.

        An explicitly passed ``adaptation`` (an ``Adaptor`` or a ``Warmup``)
        always wins; ``adaptation=None`` uses the kernel's attached default
        (``kernel.default_adaptation``). ``adaptation_kwargs`` are forwarded
        to the default factory.
        """
        if state is None or params is None or num_steps is None:
            raise TypeError(
                "warmup() missing required arguments: state, params, num_steps"
            )
        if adaptation is None:
            default = self.kernel.default_adaptation
            if default is None:
                raise ValueError(
                    "This kernel has no default adaptation; "
                    "pass adaptation= explicitly."
                )
            adaptation = default(**adaptation_kwargs)
        elif adaptation_kwargs:
            raise ValueError(
                "adaptation_kwargs need adaptation=None (they configure the default)."
            )
        return self.warmup_kernel(
            key, self.kernel, adaptation, state, params, num_steps
        )

run_kernel `staticmethod` ¶

run_kernel(key, kernel, state, num_steps, params, args=None, *, collect_state=(), collect_info=())

Run a compiled kernel scan without constructing a runner.

Source code in probjax/inference/mcmc_runner.py

@staticmethod
@partial(
    jax.jit,
    static_argnames=(
        "kernel",
        "num_steps",
        "collect_state",
        "collect_info",
    ),
)
def run_kernel(
    key: RngKey,
    kernel: MarkovKernel,
    state: State,
    num_steps: int,
    params: Params,
    args: Optional[Tuple] = None,
    *,
    collect_state: Tuple[str, ...] = (),
    collect_info: Tuple[str, ...] = (),
) -> MCMCResult:
    """Run a compiled kernel scan without constructing a runner."""
    keys = jax.random.split(key, num_steps)
    xs = keys if args is None else (keys, args)

    one_step = _mcmc_run_body(kernel, params, args, collect_state, collect_info)

    state, (states, info) = jax.lax.scan(one_step, state, xs)
    trace = states if collect_state else None
    diagnostics = info if collect_info else None
    return MCMCResult(state, trace, diagnostics)

run ¶

run(key, state, num_steps, params=None, args=None)

Run num_steps transitions with this runner's configuration.

Source code in probjax/inference/mcmc_runner.py

def run(
    self,
    key: RngKey,
    state: State,
    num_steps: int,
    params: Optional[Params] = None,
    args: Optional[Tuple] = None,
) -> MCMCResult:
    """Run ``num_steps`` transitions with this runner's configuration."""
    params = self._prepare_params(state, params)
    if not self.verbose:
        return self.run_kernel(
            key,
            self.kernel,
            state,
            num_steps,
            params,
            args,
            collect_state=self.collect_state,
            collect_info=self.collect_info,
        )

    keys = jax.random.split(key, num_steps)
    xs = keys if args is None else (keys, args)

    one_step = _mcmc_run_body(
        self.kernel,
        params,
        args,
        self.collect_state,
        self.collect_info,
        stats_fn=self._extract_stats,
    )

    state, (_, states, info) = self._verbose_scan(
        one_step,
        state,
        xs,
        num_steps,
        stats_fn=lambda _carry, y: y[0],
    )
    trace = states if self.collect_state else None
    diagnostics = info if self.collect_info else None
    return MCMCResult(state, trace, diagnostics)

sample ¶

sample(key, state, num_samples, params=None, thin=1, args=None)

Collect positions after every thin transitions.

Source code in probjax/inference/mcmc_runner.py

def sample(
    self,
    key: RngKey,
    state: State,
    num_samples: int,
    params: Optional[Params] = None,
    thin: int = 1,
    args: Optional[Tuple] = None,
) -> MCMCResult:
    """Collect positions after every ``thin`` transitions."""
    params = self._prepare_params(state, params)
    if not self.verbose:
        return self.sample_kernel(
            key,
            self.kernel,
            state,
            num_samples,
            params,
            args,
            thin=thin,
            collect_info=self.collect_info,
        )

    keys = jax.random.split(key, (num_samples, thin))
    step_args = _reshape_thin_args(args, num_samples, thin)
    xs = keys if step_args is None else (keys, step_args)

    collect_one = _mcmc_collect_body(
        self.kernel,
        params,
        step_args,
        self.collect_info,
        stats_fn=self._extract_stats,
    )

    state, (_, samples, info) = self._verbose_scan(
        collect_one,
        state,
        xs,
        num_samples,
        stats_fn=lambda _carry, y: y[0],
    )
    return MCMCResult(state, samples, info if self.collect_info else None)

adapt ¶

adapt(key, adaptor, state, params, num_steps, args=None, *, collect=False)

Fit kernel parameters with a composable adaptor.

Source code in probjax/inference/mcmc_runner.py

def adapt(
    self,
    key: RngKey,
    adaptor: Adaptor,
    state: State,
    params: Params,
    num_steps: int,
    args: Optional[Tuple] = None,
    *,
    collect: bool = False,
) -> AdaptationResult:
    """Fit kernel parameters with a composable adaptor."""
    return self.adaptor_warmup(
        key,
        self.kernel,
        adaptor,
        state,
        params,
        num_steps,
        args,
        collect=collect,
    )

warmup ¶

warmup(key, adaptation=None, state=None, params=None, num_steps=None, **adaptation_kwargs)

Run a specialized state-producing warmup procedure.

An explicitly passed adaptation (an Adaptor or a Warmup) always wins; adaptation=None uses the kernel's attached default (kernel.default_adaptation). adaptation_kwargs are forwarded to the default factory.

Source code in probjax/inference/mcmc_runner.py

def warmup(
    self,
    key: RngKey,
    adaptation: Union[Adaptor, Warmup, None] = None,
    state: Optional[State] = None,
    params: Optional[Params] = None,
    num_steps: Optional[int] = None,
    **adaptation_kwargs,
) -> WarmupResult:
    """Run a specialized state-producing warmup procedure.

    An explicitly passed ``adaptation`` (an ``Adaptor`` or a ``Warmup``)
    always wins; ``adaptation=None`` uses the kernel's attached default
    (``kernel.default_adaptation``). ``adaptation_kwargs`` are forwarded
    to the default factory.
    """
    if state is None or params is None or num_steps is None:
        raise TypeError(
            "warmup() missing required arguments: state, params, num_steps"
        )
    if adaptation is None:
        default = self.kernel.default_adaptation
        if default is None:
            raise ValueError(
                "This kernel has no default adaptation; "
                "pass adaptation= explicitly."
            )
        adaptation = default(**adaptation_kwargs)
    elif adaptation_kwargs:
        raise ValueError(
            "adaptation_kwargs need adaptation=None (they configure the default)."
        )
    return self.warmup_kernel(
        key, self.kernel, adaptation, state, params, num_steps
    )

probjax.inference.SMC ¶

Bases: WithProgressBarAPI, RunnerMixin

Compiled SMC runners, retaining evidence and final diagnostics by default.

info remains the optional legacy trace. final_info is the last kernel diagnostic, log_evidence the accumulated estimate. For a resumed fixed run pass initial_log_evidence from the earlier result (persistent states already carry it). Custom kernels without evidence diagnostics return NaN.

Source code in probjax/inference/smc_runner.py

class SMC(WithProgressBarAPI, RunnerMixin):
    """Compiled SMC runners, retaining evidence and final diagnostics by default.

    ``info`` remains the optional legacy trace. ``final_info`` is the last kernel
    diagnostic, ``log_evidence`` the accumulated estimate. For a resumed fixed
    run pass initial_log_evidence from the earlier result (persistent states
    already carry it). Custom kernels without evidence diagnostics return NaN.
    """

    _default_tracked_stats = ("ess", "log_likelihood_increment", "acceptance_rate")
    _computed_stats = {"ess": _ess_from_weights}
    _print_rate = 50

    def __init__(
        self,
        kernel,
        verbose=False,
        tracked_stats: Optional[Tuple[str, ...]] = None,
        collect=False,
    ):
        self.kernel, self.verbose, self.collect = kernel, verbose, collect
        self.tracked_stats = (
            self._default_tracked_stats if tracked_stats is None else tracked_stats
        )

    def _stat_objects(self, state, info):
        return (_sampler_state(state), info, getattr(info, "update_info", None))

    @staticmethod
    @partial(jax.jit, static_argnames=("kernel", "collect"))
    def run_kernel(
        key,
        kernel,
        state,
        tempering_params,
        params,
        *,
        collect=False,
        initial_log_evidence=None,
    ):
        return _run_fixed(
            key,
            kernel,
            state,
            tempering_params,
            params,
            collect=collect,
            initial_log_evidence=initial_log_evidence,
        )

    @staticmethod
    @partial(jax.jit, static_argnames=("kernel", "adaptor", "collect"))
    def adapt_kernel(
        key,
        kernel,
        adaptor,
        state,
        tempering_params,
        params,
        *,
        collect=False,
        initial_log_evidence=None,
    ):
        return _run_fixed(
            key,
            kernel,
            state,
            tempering_params,
            params,
            collect=collect,
            adaptor=adaptor,
            initial_log_evidence=initial_log_evidence,
        )

    @staticmethod
    @partial(jax.jit, static_argnames=("kernel", "max_steps", "adaptor"))
    def run_adaptive_kernel(
        key,
        kernel,
        state,
        params,
        *,
        max_steps=200,
        adaptor=None,
        initial_log_evidence=None,
    ):
        """Advance an adaptive kernel to temperature 1, with a fixed iteration cap.

        No history is allocated. completed=False signals the cap, exhausted
        persistent storage, nonfinite temperature/evidence, or stalled progress.
        Key/state/params can be passed back to resume a capped run.
        """
        ensure_int("max_steps", max_steps, min_value=1)
        info = _zero_info(kernel.step, key, state, mcmc_parameters=params)
        adaptation = () if adaptor is None else adaptor.init(state, params)
        evidence = _initial_evidence(state, initial_log_evidence)

        def cond(carry):
            _, state, _, _, _, _, count, valid = carry
            raw = _sampler_state(state)
            room = True
            if hasattr(raw, "persistent_log_Z"):
                room = raw.iteration + 1 < raw.persistent_log_Z.shape[0]
            return (count < max_steps) & (raw.tempering_param < 1.0) & valid & room

        def body(carry):
            key, state, params, adaptation, logz, _, count, _ = carry
            key, step_key = jax.random.split(key)
            previous = state
            state, info = kernel.step(step_key, state, mcmc_parameters=params)
            logz += _increment(previous, state, info)
            if adaptor is not None:
                adaptation, params, _ = adaptor.update(state, info, adaptation, params)
            old = _sampler_state(previous).tempering_param
            new = _sampler_state(state).tempering_param
            valid = jnp.isfinite(new) & jnp.isfinite(logz) & (new > old)
            return key, state, params, adaptation, logz, info, count + 1, valid

        key, state, params, adaptation, logz, info, count, valid = jax.lax.while_loop(
            cond,
            body,
            (
                key,
                state,
                params,
                adaptation,
                evidence,
                info,
                jnp.array(0),
                jnp.array(True),
            ),
        )
        if adaptor is not None:
            params, _ = adaptor.finalize(adaptation, params)
        params = getattr(state, "parameter_override", params)
        return SMCResult(
            state,
            params,
            None,
            logz,
            info,
            count,
            valid & (_sampler_state(state).tempering_param >= 1.0),
            key,
        )

    def run(self, key, state, tempering_params, params, *, initial_log_evidence=None):
        if self.verbose:
            return _run_fixed(
                key,
                self.kernel,
                state,
                tempering_params,
                params,
                collect=self.collect,
                verbose=self,
                initial_log_evidence=initial_log_evidence,
            )
        return self.run_kernel(
            key,
            self.kernel,
            state,
            tempering_params,
            params,
            collect=self.collect,
            initial_log_evidence=initial_log_evidence,
        )

    def sample(
        self, key, state, tempering_params, params, *, initial_log_evidence=None
    ):
        if not self.verbose:
            return self.run_kernel(
                key,
                self.kernel,
                state,
                tempering_params,
                params,
                collect=True,
                initial_log_evidence=initial_log_evidence,
            )
        return _run_fixed(
            key,
            self.kernel,
            state,
            tempering_params,
            params,
            collect=True,
            verbose=self if self.verbose else None,
            initial_log_evidence=initial_log_evidence,
        )

    def adapt(
        self,
        key,
        adaptor,
        state,
        tempering_params,
        params,
        *,
        initial_log_evidence=None,
    ):
        if not self.verbose:
            return self.adapt_kernel(
                key,
                self.kernel,
                adaptor,
                state,
                tempering_params,
                params,
                collect=self.collect,
                initial_log_evidence=initial_log_evidence,
            )
        return _run_fixed(
            key,
            self.kernel,
            state,
            tempering_params,
            params,
            collect=self.collect,
            adaptor=adaptor,
            verbose=self if self.verbose else None,
            initial_log_evidence=initial_log_evidence,
        )

    def run_adaptive(
        self,
        key,
        state,
        params,
        *,
        max_steps=200,
        adaptor=None,
        initial_log_evidence=None,
    ):
        return self.run_adaptive_kernel(
            key,
            self.kernel,
            state,
            params,
            max_steps=max_steps,
            adaptor=adaptor,
            initial_log_evidence=initial_log_evidence,
        )

run_adaptive_kernel `staticmethod` ¶

run_adaptive_kernel(key, kernel, state, params, *, max_steps=200, adaptor=None, initial_log_evidence=None)

Advance an adaptive kernel to temperature 1, with a fixed iteration cap.

No history is allocated. completed=False signals the cap, exhausted persistent storage, nonfinite temperature/evidence, or stalled progress. Key/state/params can be passed back to resume a capped run.

Source code in probjax/inference/smc_runner.py

@staticmethod
@partial(jax.jit, static_argnames=("kernel", "max_steps", "adaptor"))
def run_adaptive_kernel(
    key,
    kernel,
    state,
    params,
    *,
    max_steps=200,
    adaptor=None,
    initial_log_evidence=None,
):
    """Advance an adaptive kernel to temperature 1, with a fixed iteration cap.

    No history is allocated. completed=False signals the cap, exhausted
    persistent storage, nonfinite temperature/evidence, or stalled progress.
    Key/state/params can be passed back to resume a capped run.
    """
    ensure_int("max_steps", max_steps, min_value=1)
    info = _zero_info(kernel.step, key, state, mcmc_parameters=params)
    adaptation = () if adaptor is None else adaptor.init(state, params)
    evidence = _initial_evidence(state, initial_log_evidence)

    def cond(carry):
        _, state, _, _, _, _, count, valid = carry
        raw = _sampler_state(state)
        room = True
        if hasattr(raw, "persistent_log_Z"):
            room = raw.iteration + 1 < raw.persistent_log_Z.shape[0]
        return (count < max_steps) & (raw.tempering_param < 1.0) & valid & room

    def body(carry):
        key, state, params, adaptation, logz, _, count, _ = carry
        key, step_key = jax.random.split(key)
        previous = state
        state, info = kernel.step(step_key, state, mcmc_parameters=params)
        logz += _increment(previous, state, info)
        if adaptor is not None:
            adaptation, params, _ = adaptor.update(state, info, adaptation, params)
        old = _sampler_state(previous).tempering_param
        new = _sampler_state(state).tempering_param
        valid = jnp.isfinite(new) & jnp.isfinite(logz) & (new > old)
        return key, state, params, adaptation, logz, info, count + 1, valid

    key, state, params, adaptation, logz, info, count, valid = jax.lax.while_loop(
        cond,
        body,
        (
            key,
            state,
            params,
            adaptation,
            evidence,
            info,
            jnp.array(0),
            jnp.array(True),
        ),
    )
    if adaptor is not None:
        params, _ = adaptor.finalize(adaptation, params)
    params = getattr(state, "parameter_override", params)
    return SMCResult(
        state,
        params,
        None,
        logz,
        info,
        count,
        valid & (_sampler_state(state).tempering_param >= 1.0),
        key,
    )

MCMC kernels¶

probjax.inference.hmc ¶

Bases: MarkovKernelAPI

Hamiltonian Monte Carlo with a fixed number of integration steps.

Tunable parameters come from kernel.init_params(state, step_size=0.5, inverse_mass_matrix=None) (see :class:HMCParams); the step size is scaled by the initial gradient norm. Default adaptation: :func:window_warmup.

Parameters:

Name	Description	Default
`logdensity_fn`	Log density of the target.	required
`num_integration_steps`	Number of integrator steps per proposal.	required
`integrator`	Symplectic integrator from `blackjax.mcmc.integrators`.	required
`divergence_threshold`	Energy error above which a transition is flagged as divergent.	required

Source code in probjax/inference/mcmc/hmc.py

class hmc(MarkovKernelAPI):
    """Hamiltonian Monte Carlo with a fixed number of integration steps.

    Tunable parameters come from ``kernel.init_params(state, step_size=0.5,
    inverse_mass_matrix=None)`` (see :class:`HMCParams`); the step size is
    scaled by the initial gradient norm. Default adaptation:
    :func:`window_warmup`.

    Args:
        logdensity_fn: Log density of the target.
        num_integration_steps: Number of integrator steps per proposal.
        integrator: Symplectic integrator from ``blackjax.mcmc.integrators``.
        divergence_threshold: Energy error above which a transition is
            flagged as divergent.
    """

    init = staticmethod(blackjax.hmc.init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_step)
    default_adaptation_fn = staticmethod(window_warmup)

probjax.inference.nuts ¶

Bases: MarkovKernelAPI

No-U-Turn sampler: HMC with an adaptively chosen trajectory length.

Tunable parameters are the same as for :class:hmc (step_size, inverse_mass_matrix). Default adaptation: :func:window_warmup.

Parameters:

Name	Description	Default
`logdensity_fn`	Log density of the target.	required
`max_num_doublings`	Maximum tree depth; a trajectory has at most `2**max_num_doublings` integrator steps.	required
`divergence_threshold`	Energy error above which a transition is flagged as divergent.	required
`integrator`	Symplectic integrator from `blackjax.mcmc.integrators`.	required

Source code in probjax/inference/mcmc/hmc.py

class nuts(MarkovKernelAPI):
    """No-U-Turn sampler: HMC with an adaptively chosen trajectory length.

    Tunable parameters are the same as for :class:`hmc`
    (``step_size``, ``inverse_mass_matrix``). Default adaptation:
    :func:`window_warmup`.

    Args:
        logdensity_fn: Log density of the target.
        max_num_doublings: Maximum tree depth; a trajectory has at most
            ``2**max_num_doublings`` integrator steps.
        divergence_threshold: Energy error above which a transition is
            flagged as divergent.
        integrator: Symplectic integrator from ``blackjax.mcmc.integrators``.
    """

    init = staticmethod(blackjax.hmc.init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_kernel_nuts)
    default_adaptation_fn = staticmethod(window_warmup)

probjax.inference.dynamic_hmc ¶

Bases: MarkovKernelAPI

HMC whose number of integration steps changes from step to step.

The step count follows a deterministic or random sequence with the given average, which avoids the periodicity of fixed-length HMC. Tunable parameters are the same as for :class:hmc. Default adaptation: :func:window_warmup.

Parameters:

Name	Description	Default
`logdensity_fn`	Log density of the target.	required
`average_integration_steps`	Mean number of integrator steps.	required
`integration_steps_sequence`	How the step counts are drawn: `"halton"` (low-discrepancy sequence), `"random"` (uniform in `[1, 2 * average)`) or a `(next_random_arg_fn, integration_steps_fn)` pair.	required
`divergence_threshold`	Energy error above which a transition is flagged as divergent.	required
`integrator`	Symplectic integrator from `blackjax.mcmc.integrators`.	required

Source code in probjax/inference/mcmc/dynamic_hmc.py

class dynamic_hmc(MarkovKernelAPI):
    """HMC whose number of integration steps changes from step to step.

    The step count follows a deterministic or random sequence with the given
    average, which avoids the periodicity of fixed-length HMC. Tunable
    parameters are the same as for :class:`hmc`. Default adaptation:
    :func:`window_warmup`.

    Args:
        logdensity_fn: Log density of the target.
        average_integration_steps: Mean number of integrator steps.
        integration_steps_sequence: How the step counts are drawn:
            ``"halton"`` (low-discrepancy sequence), ``"random"`` (uniform in
            ``[1, 2 * average)``) or a ``(next_random_arg_fn,
            integration_steps_fn)`` pair.
        divergence_threshold: Energy error above which a transition is
            flagged as divergent.
        integrator: Symplectic integrator from ``blackjax.mcmc.integrators``.
    """

    init = staticmethod(init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_dynamic_hmc_step)
    default_adaptation_fn = staticmethod(window_warmup)

probjax.inference.mala ¶

Bases: MarkovKernelAPI

Metropolis-adjusted Langevin algorithm.

Proposes a gradient step plus Gaussian noise and accepts it with a Metropolis-Hastings correction. Tunable parameters come from kernel.init_params(state, step_size=1e-2) (see :class:MALAParams); the step size is scaled by the initial gradient norm. Default adaptation: :func:step_size_adaptor.

Parameters:

Name	Type	Description	Default
`logdensity_fn`		Log density of the target.	required

Source code in probjax/inference/mcmc/mala.py

class mala(MarkovKernelAPI):
    """Metropolis-adjusted Langevin algorithm.

    Proposes a gradient step plus Gaussian noise and accepts it with a
    Metropolis-Hastings correction. Tunable parameters come from
    ``kernel.init_params(state, step_size=1e-2)`` (see :class:`MALAParams`);
    the step size is scaled by the initial gradient norm. Default adaptation:
    :func:`step_size_adaptor`.

    Args:
        logdensity_fn: Log density of the target.
    """

    init = staticmethod(blackjax.mala.init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_step)
    default_adaptation_fn = staticmethod(step_size_adaptor)

probjax.inference.mclmc ¶

Bases: MarkovKernelAPI

Microcanonical Langevin Monte Carlo without Metropolis adjustment.

Isokinetic dynamics with partial momentum refreshment; asymptotically biased, so tune step_size and L with :func:mclmc_warmup (the default adaptation). Tunable parameters come from kernel.init_params(state, step_size=0.1, L=1.0, inverse_mass_matrix=None) (see :class:MCLMCParams).

Parameters:

Name	Type	Description	Default
`logdensity_fn`		Log density of the target.	required
`integrator`		Isokinetic integrator from `blackjax.mcmc.integrators`.	required

Source code in probjax/inference/mcmc/mclmc.py

class mclmc(MarkovKernelAPI):
    """Microcanonical Langevin Monte Carlo without Metropolis adjustment.

    Isokinetic dynamics with partial momentum refreshment; asymptotically
    biased, so tune ``step_size`` and ``L`` with :func:`mclmc_warmup` (the
    default adaptation). Tunable parameters come from
    ``kernel.init_params(state, step_size=0.1, L=1.0,
    inverse_mass_matrix=None)`` (see :class:`MCLMCParams`).

    Args:
        logdensity_fn: Log density of the target.
        integrator: Isokinetic integrator from ``blackjax.mcmc.integrators``.
    """

    init = staticmethod(blackjax.mclmc.init)
    init_params = staticmethod(init_mclmc_params)
    build_step = staticmethod(build_mclmc_step)
    default_adaptation_fn = staticmethod(_mclmc_adaptation)

probjax.inference.adjusted_mclmc ¶

Bases: MarkovKernelAPI

Metropolis-adjusted MCLMC with a fixed number of integration steps.

Tunable parameters come from kernel.init_params(state, step_size=0.5, num_integration_steps=10, inverse_mass_matrix=None, L_proposal_factor=inf) (see :class:AdjustedMCLMCParams). Default adaptation: :func:mclmc_warmup with adjusted=True.

Parameters:

Name	Description	Default
`logdensity_fn`	Log density of the target.	required
`integrator`	Isokinetic integrator from `blackjax.mcmc.integrators`.	required
`divergence_threshold`	Energy error above which a transition is flagged as divergent.	required

Source code in probjax/inference/mcmc/mclmc.py

class adjusted_mclmc(MarkovKernelAPI):
    """Metropolis-adjusted MCLMC with a fixed number of integration steps.

    Tunable parameters come from ``kernel.init_params(state, step_size=0.5,
    num_integration_steps=10, inverse_mass_matrix=None,
    L_proposal_factor=inf)`` (see :class:`AdjustedMCLMCParams`). Default
    adaptation: :func:`mclmc_warmup` with ``adjusted=True``.

    Args:
        logdensity_fn: Log density of the target.
        integrator: Isokinetic integrator from ``blackjax.mcmc.integrators``.
        divergence_threshold: Energy error above which a transition is
            flagged as divergent.
    """

    init = staticmethod(blackjax.adjusted_mclmc.init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_step)
    default_adaptation_fn = staticmethod(_adjusted_mclmc_adaptation)

probjax.inference.mh ¶

Bases: MarkovKernelAPI

Metropolis-Hastings with a user-supplied proposal.

The kernel params are passed to the proposal as params=; the default init_params returns the empty :class:RWParams.

Parameters:

Name	Description	Default
`logdensity_fn`	Log density of the target.	required
`transition_proposal_fn`	`(key, position, params) -> new_position`.	required
`transition_proposal_logpdf`	`(new_state, state, params) -> log q`, evaluated on BlackJAX `RWState` s; needed only for asymmetric proposals.	required

Source code in probjax/inference/mcmc/mh.py

class mh(MarkovKernelAPI):
    r"""Metropolis-Hastings with a user-supplied proposal.

    The kernel params are passed to the proposal as ``params=``; the default
    ``init_params`` returns the empty :class:`RWParams`.

    Args:
        logdensity_fn: Log density of the target.
        transition_proposal_fn: ``(key, position, params) -> new_position``.
        transition_proposal_logpdf: ``(new_state, state, params) -> log q``,
            evaluated on BlackJAX ``RWState``\ s; needed only for asymmetric
            proposals.
    """

    init = staticmethod(blackjax.rmh.init)
    init_params = staticmethod(init_params_mh)
    build_step = staticmethod(build_mh_step)

probjax.inference.gauss_rwmh ¶

Bases: MarkovKernelAPI

Gaussian random-walk Metropolis-Hastings.

Proposes x + step_size * scale @ eps with standard normal eps. Tunable parameters come from kernel.init_params(state, step_size=0.5, scale=None) (see :class:RWParamsGauss); scale is a diagonal (vector) or a dense matrix square root and defaults to the identity. Default adaptation: step size towards a 0.234 acceptance rate plus a covariance estimate for scale.

Parameters:

Name	Type	Description	Default
`logdensity_fn`		Log density of the target.	required

Source code in probjax/inference/mcmc/mh.py

class gauss_rwmh(MarkovKernelAPI):
    """Gaussian random-walk Metropolis-Hastings.

    Proposes ``x + step_size * scale @ eps`` with standard normal ``eps``.
    Tunable parameters come from ``kernel.init_params(state, step_size=0.5,
    scale=None)`` (see :class:`RWParamsGauss`); ``scale`` is a diagonal
    (vector) or a dense matrix square root and defaults to the identity.
    Default adaptation: step size towards a 0.234 acceptance rate plus a
    covariance estimate for ``scale``.

    Args:
        logdensity_fn: Log density of the target.
    """

    init = staticmethod(blackjax.rmh.init)
    init_params = staticmethod(init_params_gaussian_rw)
    build_step = staticmethod(
        partial(build_mh_step, transition_proposal_fn=gaussian_transition_proposal)
    )
    default_adaptation_fn = staticmethod(_gauss_rwmh_adaptation)

probjax.inference.imh ¶

Bases: MarkovKernelAPI

Independent Metropolis-Hastings with a user-supplied proposal.

Proposals ignore the current position. The kernel params are passed to both proposal functions as params=; the default init_params returns the empty :class:IMHParams.

Parameters:

Name	Description	Default
`logdensity_fn`	Log density of the target.	required
`proposal_fn`	`(key, params) -> position`.	required
`proposal_logpdf`	`(state, params) -> log q(state.position)`, evaluated on BlackJAX `RWState` s.	required

Source code in probjax/inference/mcmc/imh.py

class imh(MarkovKernelAPI):
    r"""Independent Metropolis-Hastings with a user-supplied proposal.

    Proposals ignore the current position. The kernel params are passed to
    both proposal functions as ``params=``; the default ``init_params``
    returns the empty :class:`IMHParams`.

    Args:
        logdensity_fn: Log density of the target.
        proposal_fn: ``(key, params) -> position``.
        proposal_logpdf: ``(state, params) -> log q(state.position)``,
            evaluated on BlackJAX ``RWState``\ s.
    """

    init = staticmethod(blackjax.irmh.init)
    init_params = staticmethod(init_imh_params)
    build_step = staticmethod(build_imh_step)

probjax.inference.neural_imh ¶

neural_imh(logdensity_fn, proposal, *, event_spec=None)

Build IMH with a tractable generative model as its proposal.

Proposal weights are explicit kernel parameters so updates made during warmup remain visible inside compiled MCMC scans.

Source code in probjax/inference/mcmc/imh.py

def neural_imh(logdensity_fn: Callable, proposal, *, event_spec=None) -> Kernel:
    """Build IMH with a tractable generative model as its proposal.

    Proposal weights are explicit kernel parameters so updates made during
    warmup remain visible inside compiled MCMC scans.
    """
    if event_spec is None:
        try:
            distribution = proposal.as_dist()
        except TypeError as error:
            raise TypeError(
                "event_spec is required for proposal models without an intrinsic "
                "event shape."
            ) from error
    else:
        distribution = proposal.as_dist(event_spec)
    if not distribution.has_logpdf:
        raise ValueError("The proposal model must expose a tractable logpdf.")
    distribution.compile("sample", "logpdf")

    def init(key, position=None, **kwargs):
        if position is None:
            position, key = key, kwargs.pop("rng_key", None)
        proposal_state = distribution.model_state()
        return NeuralIMHState(
            position,
            logdensity_fn(position),
            distribution.logpdf_with_state(proposal_state, position),
            jnp.array(0, dtype=jnp.uint32),
        )

    def init_params(state):
        del state
        return NeuralIMHParams(
            distribution.model_state(), jnp.array(0, dtype=jnp.uint32)
        )

    def step(key, state, params):
        key_proposal, key_accept = jax.random.split(key)
        proposed_position = distribution.sample_with_state(
            params.proposal_state, key_proposal
        )
        proposed_position = jax.tree.map(
            lambda proposed, current: proposed.astype(current.dtype),
            proposed_position,
            state.position,
        )
        proposed_state = NeuralIMHState(
            proposed_position,
            logdensity_fn(proposed_position),
            distribution.logpdf_with_state(params.proposal_state, proposed_position),
            params.iteration,
        )
        current_proposal_logdensity = jax.lax.cond(
            state.proposal_iteration == params.iteration,
            lambda: state.proposal_logdensity,
            lambda: distribution.logpdf_with_state(
                params.proposal_state, state.position
            ),
        )
        log_acceptance_ratio = (
            proposed_state.logdensity
            - proposed_state.proposal_logdensity
            - state.logdensity
            + current_proposal_logdensity
        )
        acceptance_rate = jnp.exp(jnp.minimum(log_acceptance_ratio, 0.0))
        is_accepted = jax.random.uniform(key_accept) < acceptance_rate
        current_state = NeuralIMHState(
            state.position,
            state.logdensity,
            current_proposal_logdensity,
            params.iteration,
        )
        next_state = jax.lax.cond(
            is_accepted, lambda: proposed_state, lambda: current_state
        )
        return next_state, RWInfo(acceptance_rate, is_accepted, proposed_state)

    return Kernel(
        init, step, init_params, lambda: neural_imh_warmup(proposal)
    )

probjax.inference.slice ¶

Bases: MarkovKernelAPI

Slice sampling along random one-dimensional slices.

Each step brackets the slice by stepping out with step_size and then shrinks it until a point is accepted. Tunable parameters come from kernel.init_params(state, step_size=0.5) (see :class:SliceParams). Default adaptation: step size towards a target number of density evaluations.

Parameters:

Name	Description	Default
`logdensity_fn`	Log density of the target.	required
`max_evals`	Maximum number of density evaluations per step.	required
`slice_fn`	Slice geometry: `"linear"` (random direction), `"poly"` (random polynomial curve), `"axis"` (random coordinate), `"none"` or a custom `(position, arg) -> (t -> position)` builder.	required
`slice_fn_arg`	`(key, position) -> arg` sampling the slice argument; inferred when `slice_fn` is a string.	required
`**kwargs`	Options for the `"poly"` slice-argument sampler.	required

Source code in probjax/inference/mcmc/slice.py

class slice(MarkovKernelAPI):
    """Slice sampling along random one-dimensional slices.

    Each step brackets the slice by stepping out with ``step_size`` and then
    shrinks it until a point is accepted. Tunable parameters come from
    ``kernel.init_params(state, step_size=0.5)`` (see :class:`SliceParams`).
    Default adaptation: step size towards a target number of density
    evaluations.

    Args:
        logdensity_fn: Log density of the target.
        max_evals: Maximum number of density evaluations per step.
        slice_fn: Slice geometry: ``"linear"`` (random direction), ``"poly"``
            (random polynomial curve), ``"axis"`` (random coordinate),
            ``"none"`` or a custom ``(position, arg) -> (t -> position)``
            builder.
        slice_fn_arg: ``(key, position) -> arg`` sampling the slice argument;
            inferred when ``slice_fn`` is a string.
        **kwargs: Options for the ``"poly"`` slice-argument sampler.
    """

    init = staticmethod(init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_step)
    default_adaptation_fn = staticmethod(_slice_adaptation)

probjax.inference.latent_slice ¶

Bases: MarkovKernelAPI

Latent slice sampling with an auxiliary, per-coordinate bracket width.

The bracket width is resampled every step from an exponential with rate step_size, so it adapts to the local scale. Tunable parameters come from kernel.init_params(state, step_size=0.5) (see :class:LatentSliceParams). Default adaptation: step size towards a target number of density evaluations.

Parameters:

Name	Type	Description	Default
`logdensity_fn`		Log density of the target.	required
`max_evals`		Maximum number of density evaluations per step.	required

Source code in probjax/inference/mcmc/latent_slice.py

class latent_slice(MarkovKernelAPI):
    """Latent slice sampling with an auxiliary, per-coordinate bracket width.

    The bracket width is resampled every step from an exponential with rate
    ``step_size``, so it adapts to the local scale. Tunable parameters come
    from ``kernel.init_params(state, step_size=0.5)`` (see
    :class:`LatentSliceParams`). Default adaptation: step size towards a
    target number of density evaluations.

    Args:
        logdensity_fn: Log density of the target.
        max_evals: Maximum number of density evaluations per step.
    """

    init = staticmethod(init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_step)
    default_adaptation_fn = staticmethod(_latent_slice_adaptation)

probjax.inference.elliptical_slice ¶

Bases: MarkovKernelAPI

Elliptical slice sampling for targets with a Gaussian prior.

The target is split as N(mean, cov_matrix) * likelihood. Rejection free and without tunable parameters.

Parameters:

Name	Description	Default
`logdensity_fn`	Log density of the full target.	required
`cov_matrix`	Prior covariance, a diagonal (vector) or a dense matrix.	required
`mean`	Prior mean.	required
`loglikelihood_fn`	Log likelihood. Defaults to `logdensity_fn` minus the Gaussian prior log density.	required

Source code in probjax/inference/mcmc/elliptical_slice.py

class elliptical_slice(MarkovKernelAPI):
    """Elliptical slice sampling for targets with a Gaussian prior.

    The target is split as ``N(mean, cov_matrix) * likelihood``. Rejection
    free and without tunable parameters.

    Args:
        logdensity_fn: Log density of the full target.
        cov_matrix: Prior covariance, a diagonal (vector) or a dense matrix.
        mean: Prior mean.
        loglikelihood_fn: Log likelihood. Defaults to ``logdensity_fn`` minus
            the Gaussian prior log density.
    """

    init = staticmethod(blackjax.elliptical_slice.init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_elliptical_slice_step)

probjax.inference.arms ¶

Bases: MarkovKernelAPI

Adaptive rejection Metropolis sampling for scalar targets.

Proposes from a piecewise-exponential envelope of the log density and corrects with Metropolis-Hastings; rejected proposals refine the envelope. The envelope is set up by kernel.init(key, position, init_points=None, lb=-inf, ub=inf, max_points=50). No tunable parameters.

Parameters:

Name	Type	Description	Default
`logdensity_fn`		Log density of the scalar target.	required

Source code in probjax/inference/mcmc/arms.py

class arms(MarkovKernelAPI):
    """Adaptive rejection Metropolis sampling for scalar targets.

    Proposes from a piecewise-exponential envelope of the log density and
    corrects with Metropolis-Hastings; rejected proposals refine the
    envelope. The envelope is set up by ``kernel.init(key, position,
    init_points=None, lb=-inf, ub=inf, max_points=50)``. No tunable
    parameters.

    Args:
        logdensity_fn: Log density of the scalar target.
    """

    init = staticmethod(init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(
        lambda logdensity_fn, **_: _arms_step(logdensity_fn, control=False)
    )

probjax.inference.a2rms ¶

Bases: MarkovKernelAPI

ARMS that also refines the envelope after some accepted proposals.

Like :class:arms, but an accepted proposal still updates the envelope with a probability that grows with the envelope's gap to the target, so the envelope keeps improving where ARMS would stop.

Parameters:

Name	Type	Description	Default
`logdensity_fn`		Log density of the scalar target.	required

Source code in probjax/inference/mcmc/arms.py

class a2rms(MarkovKernelAPI):
    """ARMS that also refines the envelope after some accepted proposals.

    Like :class:`arms`, but an accepted proposal still updates the envelope
    with a probability that grows with the envelope's gap to the target, so
    the envelope keeps improving where ARMS would stop.

    Args:
        logdensity_fn: Log density of the scalar target.
    """

    init = staticmethod(init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(
        lambda logdensity_fn, **_: _arms_step(logdensity_fn, control=True)
    )

probjax.inference.ars ¶

ars(state, key, num_samples, log_density_fn)

Source code in probjax/inference/rejection/univariate.py

@partial(jax.jit, static_argnums=(2, 3))
def ars(state: ARSState, key: RngKey, num_samples: int, log_density_fn: Callable):
    samples = jnp.zeros(num_samples)
    n = 0

    def cond_fn(carry):
        _, n, _, _, _ = carry
        return n < num_samples

    def body_fn(carry):
        key, n, state, samples, iteration = carry
        key, sample_key, rejection_key = jax.random.split(key, 3)

        [xt, i] = sample_upper(state, sample_key)

        lh = eval_lower(state, xt)
        uh = eval_upper(state, xt, i)

        u = jax.random.uniform(rejection_key)

        def accept(xt, n, state, samples):
            samples = samples.at[n].set(xt)
            n += 1
            return samples, n, state

        def reject(xt, n, state, samples):
            return samples, n, state

        def reject_maybe_accept_update(xt, n, state, samples):
            h = log_density_fn(xt)
            hprime = jax.grad(log_density_fn)(xt)

            samples, n, state = jax.lax.cond(
                u <= jnp.exp(h - uh), accept, reject, xt, n, state, samples
            )
            state = update_ars_state(state, xt, h, hprime)
            return samples, n, state

        samples, n, state = jax.lax.cond(
            (u <= jnp.exp(lh - uh)) & (lh <= uh),
            accept,
            reject_maybe_accept_update,
            xt,
            n,
            state,
            samples,
        )

        return (key, n, state, samples, iteration + 1)

    carry = (key, n, state, samples, 0)
    _, _, state, samples, iterations = jax.lax.while_loop(cond_fn, body_fn, carry)

    return samples, state, num_samples / iterations

probjax.inference.gaussian_imh ¶

Bases: MarkovKernelAPI

Independent Metropolis-Hastings with a Gaussian proposal.

Tunable parameters come from kernel.init_params(state, mean=None, cov=None) (see :class:GaussianIMHParams); mean defaults to the initial position and cov (a diagonal vector or a dense matrix) to the identity. Default adaptation: running estimates of the mean and covariance.

Parameters:

Name	Type	Description	Default
`logdensity_fn`		Log density of the target.	required

Source code in probjax/inference/mcmc/imh.py

class gaussian_imh(MarkovKernelAPI):
    """Independent Metropolis-Hastings with a Gaussian proposal.

    Tunable parameters come from ``kernel.init_params(state, mean=None,
    cov=None)`` (see :class:`GaussianIMHParams`); ``mean`` defaults to the
    initial position and ``cov`` (a diagonal vector or a dense matrix) to the
    identity. Default adaptation: running estimates of the mean and
    covariance.

    Args:
        logdensity_fn: Log density of the target.
    """

    init = staticmethod(blackjax.irmh.init)
    init_params = staticmethod(init_gaussian_imh_params)
    build_step = staticmethod(build_gaussian_imh_step)
    default_adaptation_fn = staticmethod(_gaussian_imh_adaptation)

probjax.inference.adjusted_mclmc_dynamic ¶

Bases: MarkovKernelAPI

Metropolis-adjusted MCLMC with a varying number of integration steps.

Tunable parameters come from kernel.init_params(state, step_size=0.5, inverse_mass_matrix=None, L_proposal_factor=inf) (see :class:AdjustedMCLMCDynamicParams). Default adaptation: :func:mclmc_warmup with adjusted=True.

Parameters:

Name	Description	Default
`logdensity_fn`	Log density of the target.	required
`average_integration_steps`	Mean number of integrator steps.	required
`integration_steps_sequence`	How the step counts are drawn: `"halton"`, `"random"` or a `(next_random_arg_fn, integration_steps_fn)` pair.	required
`divergence_threshold`	Energy error above which a transition is flagged as divergent.	required
`integrator`	Isokinetic integrator from `blackjax.mcmc.integrators`.	required
`next_random_arg_fn`	Advances the random argument; required with `integration_steps_fn`.	required
`integration_steps_fn`	Custom step-count function, overriding `integration_steps_sequence`.	required

Source code in probjax/inference/mcmc/mclmc.py

class adjusted_mclmc_dynamic(MarkovKernelAPI):
    """Metropolis-adjusted MCLMC with a varying number of integration steps.

    Tunable parameters come from ``kernel.init_params(state, step_size=0.5,
    inverse_mass_matrix=None, L_proposal_factor=inf)`` (see
    :class:`AdjustedMCLMCDynamicParams`). Default adaptation:
    :func:`mclmc_warmup` with ``adjusted=True``.

    Args:
        logdensity_fn: Log density of the target.
        average_integration_steps: Mean number of integrator steps.
        integration_steps_sequence: How the step counts are drawn:
            ``"halton"``, ``"random"`` or a ``(next_random_arg_fn,
            integration_steps_fn)`` pair.
        divergence_threshold: Energy error above which a transition is
            flagged as divergent.
        integrator: Isokinetic integrator from ``blackjax.mcmc.integrators``.
        next_random_arg_fn: Advances the random argument; required with
            ``integration_steps_fn``.
        integration_steps_fn: Custom step-count function, overriding
            ``integration_steps_sequence``.
    """

    init = staticmethod(init_dynamic)
    init_params = staticmethod(init_dynamic_params)
    build_step = staticmethod(build_dynamic_step)
    default_adaptation_fn = staticmethod(_adjusted_mclmc_adaptation)

probjax.inference.RejectionSampler ¶

Source code in probjax/inference/rejection/__init__.py

class RejectionSampler:
    def __init__(
        self,
        potential_fn: Callable,
        proposal: Distribution,
        trial_samples_logM: int = 1000,
        batch_size: int = 100,
    ) -> None:
        self._potential_fn = potential_fn
        self._proposal = proposal
        self._log_density_ratio_fn = lambda x: (
            self._potential_fn(x) - self._proposal.log_prob(x)
        )
        self._log_M = estimate_ratio_bound(
            self._log_density_ratio_fn,
            proposal.sample(jax.random.PRNGKey(42), (trial_samples_logM,)),
        )
        self._batch_size = batch_size

    def run(self, key, num_samples: int = 1, **kwargs) -> Array:
        key, subkey = jax.random.split(key)
        batch_size = self._batch_size
        samples = jnp.empty((num_samples,) + self._proposal.event_shape)

        def cond_fn(state):
            i, key, samples = state
            return i < num_samples

        def body_fn(state):
            i, key, samples = state
            key, key_propose, key_accept = jax.random.split(key, 3)
            proposed_samples = self._proposal.sample(
                key_propose, (batch_size,), **kwargs
            )
            log_ratio = self._log_density_ratio_fn(proposed_samples)
            log_acceptance = log_ratio - self._log_M
            accept = log_acceptance > jax.random.uniform(
                key_accept, shape=log_acceptance.shape
            )
            num_accepted = jnp.sum(accept)
            samples = samples[accept].set(proposed_samples[accept])
            i = i + num_accepted
            return i, key, samples

        _, _, samples = jax.lax.while_loop(cond_fn, body_fn, (0, key, samples))
        return samples

probjax.inference.pseudo_marginal ¶

pseudo_marginal(inner_kernel_cls, stochastic_logdensity_fn, num_samples=1, **inner_kernel_kwargs)

Wrap a probjax MCMC kernel for pseudo-marginal inference.

Parameters:

Name	Type	Description	Default
`inner_kernel_cls`		A probjax kernel class (e.g. `hmc`, `mala`, `gauss_rwmh`) that exposes `build_step`, `init`, and `init_params`.	required
`stochastic_logdensity_fn`	`Callable`	A callable `(position, rng_key) -> float` that returns an unbiased estimate of the log-density (or log-marginal-likelihood) given auxiliary randomness from rng_key.	required
`num_samples`	`int`	Number of independent keys to evaluate and average for variance reduction (default 1).	`1`
`**inner_kernel_kwargs`		Forwarded to `inner_kernel_cls.build_step` (e.g. `num_integration_steps` for HMC).	`{}`

Returns:

Name	Type	Description
`A`	`MarkovKernel`	class:`MarkovKernel` whose `step` splits the PRNG key,
	`MarkovKernel`	fixes one sub-key for the stochastic log-density, and delegates
	`MarkovKernel`	the transition to the inner kernel.

Source code in probjax/inference/mcmc/pmmcmc.py

def pseudo_marginal(
    inner_kernel_cls,
    stochastic_logdensity_fn: Callable,
    num_samples: int = 1,
    **inner_kernel_kwargs,
) -> MarkovKernel:
    """Wrap a probjax MCMC kernel for pseudo-marginal inference.

    Args:
        inner_kernel_cls: A probjax kernel class (e.g. ``hmc``, ``mala``,
            ``gauss_rwmh``) that exposes ``build_step``, ``init``, and
            ``init_params``.
        stochastic_logdensity_fn: A callable
            ``(position, rng_key) -> float`` that returns an unbiased
            estimate of the log-density (or log-marginal-likelihood)
            given auxiliary randomness from *rng_key*.
        num_samples: Number of independent keys to evaluate and average
            for variance reduction (default 1).
        **inner_kernel_kwargs: Forwarded to
            ``inner_kernel_cls.build_step`` (e.g. ``num_integration_steps``
            for HMC).

    Returns:
        A :class:`MarkovKernel` whose ``step`` splits the PRNG key,
        fixes one sub-key for the stochastic log-density, and delegates
        the transition to the inner kernel.
    """
    if num_samples < 1:
        raise ValueError("num_samples must be >= 1")

    # ------------------------------------------------------------------
    # init: evaluate the stochastic logdensity once to populate state
    # ------------------------------------------------------------------
    def init(key, position=None, rng_key=None):
        if position is None:
            position, key = key, rng_key
        if key is None:
            key = jax.random.PRNGKey(0)
        key_init, key_logdensity = jax.random.split(key)
        fixed_logdensity = _make_fixed_logdensity(
            stochastic_logdensity_fn, key_logdensity, num_samples
        )
        # The inner kernel's underlying init (e.g. blackjax.hmc.init)
        # takes (position, logdensity_fn) and evaluates at position.
        return inner_kernel_cls.init(
            position, logdensity_fn=fixed_logdensity, rng_key=key_init
        )

    # ------------------------------------------------------------------
    # step: fix randomness for this transition, delegate to inner kernel
    # ------------------------------------------------------------------
    def step(key: RngKey, state, params, *args):
        key_kernel, key_logdensity = jax.random.split(key)
        fixed_logdensity = _make_fixed_logdensity(
            stochastic_logdensity_fn, key_logdensity, num_samples
        )
        # build_step returns a raw step(key, state, params) closure
        inner_step = inner_kernel_cls.build_step(
            fixed_logdensity, **inner_kernel_kwargs
        )
        return inner_step(key_kernel, state, params, *args)

    def init_params(state, *args, **kwargs):
        return inner_kernel_cls.init_params(state, *args, **kwargs)

    # Inherit the inner kernel's recommended warmup (None if it has none).
    return MarkovKernel(
        init, step, init_params, getattr(inner_kernel_cls, "default_adaptation_fn", None)
    )

Stochastic-gradient MCMC¶

probjax.inference.sgld ¶

sgld(grad_estimator, temperature=1.0)

Stochastic Gradient Langevin Dynamics.

Parameters:

Name	Type	Description	Default
`grad_estimator`	`Callable`	Function `(position, minibatch) -> pytree` that estimates the gradient of the log-posterior.	required
`temperature`	`float`	Temperature parameter (default 1.0).	`1.0`

Source code in probjax/inference/mcmc/sgmcmc.py

def sgld(grad_estimator: Callable, temperature: float = 1.0) -> MarkovKernel:
    """Stochastic Gradient Langevin Dynamics.

    Args:
        grad_estimator: Function ``(position, minibatch) -> pytree`` that
            estimates the gradient of the log-posterior.
        temperature: Temperature parameter (default 1.0).
    """
    kernel = _make_sgmcmc_kernel(
        grad_estimator,
        blackjax.sgld.build_kernel,
        blackjax_init=blackjax.sgld.init,
    )
    return kernel

probjax.inference.sghmc ¶

sghmc(grad_estimator, num_integration_steps=10, alpha=0.01, beta=0.0, temperature=1.0)

Stochastic Gradient Hamiltonian Monte Carlo.

Parameters:

Name	Type	Description	Default
`grad_estimator`	`Callable`	Function `(position, minibatch) -> pytree` that estimates the gradient of the log-posterior.	required
`num_integration_steps`	`int`	Number of leapfrog steps per sample.	`10`
`alpha`	`float`	Friction coefficient.	`0.01`
`beta`	`float`	Noise scaling.	`0.0`
`temperature`	`float`	Temperature parameter.	`1.0`

Source code in probjax/inference/mcmc/sgmcmc.py

def sghmc(
    grad_estimator: Callable,
    num_integration_steps: int = 10,
    alpha: float = 0.01,
    beta: float = 0.0,
    temperature: float = 1.0,
) -> MarkovKernel:
    """Stochastic Gradient Hamiltonian Monte Carlo.

    Args:
        grad_estimator: Function ``(position, minibatch) -> pytree`` that
            estimates the gradient of the log-posterior.
        num_integration_steps: Number of leapfrog steps per sample.
        alpha: Friction coefficient.
        beta: Noise scaling.
        temperature: Temperature parameter.
    """
    return _make_sgmcmc_kernel(
        grad_estimator,
        partial(blackjax.sghmc.build_kernel, alpha=alpha, beta=beta),
        blackjax_init=blackjax.sghmc.init,
        extra_step_kwargs={"num_integration_steps": num_integration_steps},
    )

probjax.inference.sgnht ¶

sgnht(grad_estimator, alpha=0.01, beta=0.0, temperature=1.0)

Stochastic Gradient Nosé-Hoover Thermostat.

SGNHT extends SGHMC with a thermostat variable that automatically adjusts the kinetic energy to maintain the target temperature.

Parameters:

Name	Type	Description	Default
`grad_estimator`	`Callable`	Function `(position, minibatch) -> pytree` that estimates the gradient of the log-posterior.	required
`alpha`	`float`	Friction coefficient.	`0.01`
`beta`	`float`	Noise scaling.	`0.0`
`temperature`	`float`	Temperature parameter.	`1.0`

Note

init requires rng_key to initialise momentum::

state = sampler.init(position, rng_key=key)

Source code in probjax/inference/mcmc/sgmcmc.py

def sgnht(
    grad_estimator: Callable,
    alpha: float = 0.01,
    beta: float = 0.0,
    temperature: float = 1.0,
) -> MarkovKernel:
    """Stochastic Gradient Nosé-Hoover Thermostat.

    SGNHT extends SGHMC with a thermostat variable that automatically
    adjusts the kinetic energy to maintain the target temperature.

    Args:
        grad_estimator: Function ``(position, minibatch) -> pytree`` that
            estimates the gradient of the log-posterior.
        alpha: Friction coefficient.
        beta: Noise scaling.
        temperature: Temperature parameter.

    Note:
        ``init`` requires ``rng_key`` to initialise momentum::

            state = sampler.init(position, rng_key=key)
    """
    kernel = blackjax.sgnht.build_kernel(alpha=alpha, beta=beta)

    def sgnht_state_init(key, position, rng_key):
        if position is None:
            position, key = key, rng_key
        return _sgnht_init(position, rng_key=key, alpha=alpha)

    return _make_sgmcmc_kernel(
        grad_estimator,
        lambda: kernel,
        state_init=sgnht_state_init,
    )

Warmup and adaptation¶

probjax.inference.window_warmup ¶

window_warmup(*, diagonal=True, target_acceptance_rate=0.8, initial_buffer_size=75, final_buffer_size=50, first_window_size=25)

Finite Stan-style warmup using BlackJAX's window-adaptation base.

Source code in probjax/inference/mcmc/adaptation.py

def window_warmup(
    *,
    diagonal: bool = True,
    target_acceptance_rate: float = 0.8,
    initial_buffer_size: int = 75,
    final_buffer_size: int = 50,
    first_window_size: int = 25,
) -> Warmup:
    """Finite Stan-style warmup using BlackJAX's window-adaptation base."""
    import warnings

    # `window_adaptation.base` is deprecated upstream in favor of the bound
    # `window_adaptation`/`staged_adaptation` algorithms, but those don't expose
    # the init/update/final triple this Warmup protocol needs. Keep using the
    # base primitives and silence the deprecation until the protocol migrates
    # to the staged engine.
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore", message=".*window_adaptation.base.*", category=DeprecationWarning
        )
        init_window, update_window, final_window = window_adaptation_base(
            diagonal, target_acceptance_rate
        )

    def run(key, kernel, state, params, num_steps, args=None):
        if args is not None:
            raise ValueError("window_warmup does not support per-step arguments")
        schedule = build_schedule(
            num_steps,
            initial_buffer_size=initial_buffer_size,
            final_buffer_size=final_buffer_size,
            first_window_size=first_window_size,
        )
        get_param(params, "inverse_mass_matrix")
        warmup_state = init_window(state.position, get_param(params, "step_size"))
        keys = jax.random.split(key, num_steps)

        def one_step(carry, xs):
            state, params, warmup_state = carry
            step_key, stage = xs
            state, info = kernel.step(step_key, state, params)
            warmup_state = update_window(
                warmup_state, stage, state.position, info.acceptance_rate
            )
            params = replace_params(
                params,
                step_size=warmup_state.step_size,
                inverse_mass_matrix=warmup_state.inverse_mass_matrix,
            )
            return (state, params, warmup_state), None

        (state, params, warmup_state), _ = jax.lax.scan(
            one_step, (state, params, warmup_state), (keys, schedule)
        )
        step_size, inverse_mass_matrix = final_window(warmup_state)
        params = replace_params(
            params,
            step_size=step_size,
            inverse_mass_matrix=inverse_mass_matrix,
        )
        return WarmupResult(state, params)

    return Warmup(run)

probjax.inference.pathfinder_warmup ¶

pathfinder_warmup(algorithm, logdensity_fn, *, initial_step_size=1.0, target_acceptance_rate=0.8, collect=False, **extra_parameters)

Create a BlackJAX Pathfinder warmup procedure.

Source code in probjax/inference/mcmc/warmup.py

def pathfinder_warmup(
    algorithm,
    logdensity_fn,
    *,
    initial_step_size: float = 1.0,
    target_acceptance_rate: float = 0.8,
    collect: bool = False,
    **extra_parameters,
) -> Warmup:
    """Create a BlackJAX Pathfinder warmup procedure."""
    info_fn = return_all_adapt_info if collect else lambda *_: None
    procedure = blackjax.pathfinder_adaptation(
        algorithm,
        logdensity_fn,
        initial_step_size=initial_step_size,
        target_acceptance_rate=target_acceptance_rate,
        adaptation_info_fn=info_fn,
        **extra_parameters,
    )

    @partial(jax.jit, static_argnames=("num_steps",))
    def run(key, _kernel, state, params, num_steps):
        result, info = procedure.run(key, state.position, num_steps)
        params = replace_params(params, **result.parameters)
        return WarmupResult(result.state, params, info if collect else None)

    return Warmup(run)

probjax.inference.mclmc_warmup ¶

mclmc_warmup(mclmc_kernel, *, adjusted=False, target_acceptance_rate=0.8, collect=False, **options)

Create a specialized BlackJAX MCLMC warmup procedure.

mclmc_kernel follows the corresponding BlackJAX adaptation kernel protocol. MCLMC warmup advances the chain while estimating L, step size, and optional diagonal preconditioning.

Source code in probjax/inference/mcmc/warmup.py

def mclmc_warmup(
    mclmc_kernel,
    *,
    adjusted: bool = False,
    target_acceptance_rate: float = 0.8,
    collect: bool = False,
    **options,
) -> Warmup:
    """Create a specialized BlackJAX MCLMC warmup procedure.

    ``mclmc_kernel`` follows the corresponding BlackJAX adaptation kernel
    protocol. MCLMC warmup advances the chain while estimating ``L``, step size,
    and optional diagonal preconditioning.
    """
    finder = (
        blackjax.adjusted_mclmc_find_L_and_step_size
        if adjusted
        else blackjax.mclmc_find_L_and_step_size
    )

    @partial(jax.jit, static_argnames=("num_steps",))
    def run(key, _kernel, state, params, num_steps):
        adaptation_params = MCLMCAdaptationState(
            get_param(params, "L"),
            get_param(params, "step_size"),
            get_param(params, "inverse_mass_matrix"),
        )
        kwargs = dict(options)
        if adjusted:
            kwargs["target"] = target_acceptance_rate
        state, adaptation_params, work = finder(
            mclmc_kernel=mclmc_kernel,
            num_steps=num_steps,
            state=state,
            rng_key=key,
            params=adaptation_params,
            **kwargs,
        )
        params = replace_params(
            params,
            L=adaptation_params.L,
            step_size=adaptation_params.step_size,
            inverse_mass_matrix=adaptation_params.inverse_mass_matrix,
        )
        return WarmupResult(state, params, work if collect else None)

    return Warmup(run)

probjax.inference.neural_imh_warmup ¶

neural_imh_warmup(proposal, *, data=None, num_adaptations=5, fit_steps=100, batch_size=256, max_buffer_size=10000, rao_blackwellize=True, learning_rate=0.001, optimizer=None)

Adapt a neural IMH proposal on chain positions and optional seed data.

The requested warmup transitions are split into num_adaptations blocks. After each block, the proposal is fitted to positions collected so far, augmented by data when supplied. By default each transition contributes its current and proposed positions with weights 1 - alpha and alpha. The returned parameters are then fixed for regular MCMC sampling.

Source code in probjax/inference/mcmc/imh.py

def neural_imh_warmup(
    proposal,
    *,
    data=None,
    num_adaptations: int = 5,
    fit_steps: int = 100,
    batch_size: Optional[int] = 256,
    max_buffer_size: Optional[int] = 10_000,
    rao_blackwellize: bool = True,
    learning_rate: float = 1e-3,
    optimizer=None,
) -> Warmup:
    """Adapt a neural IMH proposal on chain positions and optional seed data.

    The requested warmup transitions are split into ``num_adaptations`` blocks.
    After each block, the proposal is fitted to positions collected so far,
    augmented by ``data`` when supplied. By default each transition contributes
    its current and proposed positions with weights ``1 - alpha`` and ``alpha``.
    The returned parameters are then fixed for regular MCMC sampling.
    """
    if num_adaptations < 1:
        raise ValueError("num_adaptations must be at least one.")
    if fit_steps < 1:
        raise ValueError("fit_steps must be at least one.")
    if max_buffer_size is not None and max_buffer_size < 1:
        raise ValueError("max_buffer_size must be positive or None.")
    fit_optimizer = optimizer
    if fit_optimizer is None:
        fit_optimizer = optax.adam(learning_rate)

    def run(key, kernel, state, params, num_steps):
        if not isinstance(params, NeuralIMHParams):
            raise TypeError("neural_imh_warmup requires a neural_imh kernel.")
        if num_steps < num_adaptations:
            raise ValueError("num_steps must be at least num_adaptations.")

        nnx.update(proposal, params.proposal_state)

        quotient, remainder = divmod(num_steps, num_adaptations)
        block_sizes = [
            quotient + (round_index < remainder)
            for round_index in range(num_adaptations)
        ]
        collected = None
        collected_weights = None
        losses = []
        acceptance_rates = []

        for block_size in block_sizes:
            key, sample_key, fit_key = jax.random.split(key, 3)
            initial_position = state.position
            collect_info = (
                ("acceptance_rate", "proposal")
                if rao_blackwellize
                else ("acceptance_rate",)
            )
            result = sample_kernel(
                sample_key,
                kernel,
                state,
                block_size,
                params,
                collect_info=collect_info,
            )
            state = result.state
            if rao_blackwellize:
                round_data, round_weights = _rao_blackwellized_data(
                    initial_position,
                    result.samples,
                    result.info["proposal"].position,
                    result.info["acceptance_rate"],
                )
            else:
                round_data = result.samples
                round_weights = jnp.ones((block_size,))
            collected = _limit_data(
                _concatenate_data(collected, round_data), max_buffer_size
            )
            collected_weights = _limit_data(
                _concatenate_data(collected_weights, round_weights),
                max_buffer_size,
            )
            fit_data = _concatenate_data(data, collected)
            fit_weights = collected_weights
            if data is not None:
                num_seed = jax.tree.leaves(data)[0].shape[0]
                fit_weights = jnp.concatenate((jnp.ones(num_seed), fit_weights))
            round_losses = proposal.fit(
                fit_key,
                fit_data,
                weights=fit_weights,
                num_steps=fit_steps,
                batch_size=batch_size,
                learning_rate=learning_rate,
                optimizer=fit_optimizer,
            )
            _, proposal_state = nnx.split(proposal)
            proposal_state = jax.tree.map(lambda value: value.copy(), proposal_state)
            params = NeuralIMHParams(proposal_state, params.iteration + 1)
            losses.append(round_losses)
            acceptance_rates.append(jnp.mean(result.info["acceptance_rate"]))

        info = NeuralIMHWarmupInfo(
            losses=jnp.stack(losses),
            acceptance_rate=jnp.stack(acceptance_rates),
        )
        return WarmupResult(state, params, info)

    return Warmup(run)

probjax.inference.step_size_adaptor ¶

step_size_adaptor(target=0.8, *, target_from_info_fn=lambda info: acceptance_rate, t0=10, gamma=0.05, kappa=0.75)

Local dual-averaging rule backed by BlackJAX's adaptation primitive.

Source code in probjax/inference/mcmc/adaptation.py

def step_size_adaptor(
    target: float = 0.8,
    *,
    target_from_info_fn=lambda info: info.acceptance_rate,
    t0: int = 10,
    gamma: float = 0.05,
    kappa: float = 0.75,
) -> Adaptor:
    """Local dual-averaging rule backed by BlackJAX's adaptation primitive."""
    init_da, update_da, final_da = dual_averaging_adaptation(
        target, t0=t0, gamma=gamma, kappa=kappa
    )

    def init(_state, params):
        return init_da(get_param(params, "step_size"))

    def update(_state, info, adaptor_state, params):
        adaptor_state = update_da(adaptor_state, target_from_info_fn(info))
        params = replace_params(params, step_size=jnp.exp(adaptor_state.log_step_size))
        return adaptor_state, params, None

    def finalize(adaptor_state, params):
        return replace_params(params, step_size=final_da(adaptor_state)), None

    return Adaptor(init, update, finalize)

probjax.inference.mass_matrix_adaptor ¶

mass_matrix_adaptor(*, diagonal=True)

Local mass-matrix estimator backed by BlackJAX Welford adaptation.

Source code in probjax/inference/mcmc/adaptation.py

def mass_matrix_adaptor(*, diagonal: bool = True) -> Adaptor:
    """Local mass-matrix estimator backed by BlackJAX Welford adaptation."""
    init_mm, update_mm, final_mm = mass_matrix_adaptation(diagonal)

    def init(state, params):
        position, _ = jax.flatten_util.ravel_pytree(state.position)
        get_param(params, "inverse_mass_matrix")
        return init_mm(position.size)

    def update(state, _info, adaptor_state, params):
        return update_mm(adaptor_state, state.position), params, None

    def finalize(adaptor_state, params):
        adaptor_state = final_mm(adaptor_state)
        return replace_params(
            params, inverse_mass_matrix=adaptor_state.inverse_mass_matrix
        ), None

    return Adaptor(init, update, finalize)

probjax.inference.covariance_adaptor ¶

covariance_adaptor(*, field='scale', diagonal=True, output='standard_deviation')

Local proposal-geometry estimator backed by BlackJAX Welford updates.

Source code in probjax/inference/mcmc/adaptation.py

def covariance_adaptor(
    *,
    field: str = "scale",
    diagonal: bool = True,
    output: str = "standard_deviation",
) -> Adaptor:
    """Local proposal-geometry estimator backed by BlackJAX Welford updates."""
    init_cov, update_cov, final_cov = welford_algorithm(diagonal)
    if output not in {"variance", "covariance", "standard_deviation", "cholesky"}:
        raise ValueError(f"Unsupported covariance output: {output}")

    def init(state, params):
        position, _ = jax.flatten_util.ravel_pytree(state.position)
        get_param(params, field)
        return CovarianceAdaptorState(init_cov(position.size))

    def update(state, _info, adaptor_state, params):
        position, _ = jax.flatten_util.ravel_pytree(state.position)
        state_out = update_cov(adaptor_state.welford_state, position)
        return CovarianceAdaptorState(state_out), params, None

    def finalize(adaptor_state, params):
        covariance, _, _ = final_cov(adaptor_state.welford_state)
        if output == "standard_deviation":
            value = jnp.sqrt(covariance)
        elif output == "cholesky":
            value = jnp.linalg.cholesky(covariance)
        else:
            value = covariance
        return replace_params(params, **{field: value}), None

    return Adaptor(init, update, finalize)

probjax.inference.mean_adaptor ¶

mean_adaptor(*, field='mean', diagonal=True)

Local proposal-mean estimator backed by BlackJAX Welford updates.

Source code in probjax/inference/mcmc/adaptation.py

def mean_adaptor(
    *,
    field: str = "mean",
    diagonal: bool = True,
) -> Adaptor:
    """Local proposal-mean estimator backed by BlackJAX Welford updates."""
    init_mean, update_mean, final_mean = welford_algorithm(diagonal)

    def init(state, params):
        position, _ = jax.flatten_util.ravel_pytree(state.position)
        get_param(params, field)
        return CovarianceAdaptorState(init_mean(position.size))

    def update(state, _info, adaptor_state, params):
        position, _ = jax.flatten_util.ravel_pytree(state.position)
        state_out = update_mean(adaptor_state.welford_state, position)
        return CovarianceAdaptorState(state_out), params, None

    def finalize(adaptor_state, params):
        _, _, mean = final_mean(adaptor_state.welford_state)
        return replace_params(params, **{field: mean}), None

    return Adaptor(init, update, finalize)

probjax.inference.compose_adaptors ¶

compose_adaptors(**adaptors)

Compose adaptors, threading parameter updates in declaration order.

Source code in probjax/inference/adaptation.py

def compose_adaptors(**adaptors) -> Adaptor:
    """Compose adaptors, threading parameter updates in declaration order."""
    if not adaptors:
        raise ValueError("At least one adaptor is required")
    names = tuple(adaptors)
    children = tuple(adaptors.values())

    def init(state, params):
        return tuple(child.init(state, params) for child in children)

    def update(state, info, adaptor_states, params):
        new_states = []
        child_info = {}
        for name, child, child_state in zip(
            names, children, adaptor_states, strict=True
        ):
            child_state, params, info_out = child.update(
                state, info, child_state, params
            )
            new_states.append(child_state)
            child_info[name] = info_out
        return tuple(new_states), params, child_info

    def finalize(adaptor_states, params):
        child_info = {}
        for name, child, child_state in zip(
            names, children, adaptor_states, strict=True
        ):
            params, info_out = child.finalize(child_state, params)
            child_info[name] = info_out
        return params, child_info

    return Adaptor(init, update, finalize)

probjax.inference.acceptance_rate_adaptor ¶

acceptance_rate_adaptor(*, target_acceptance_rate=0.234, scale_key='step_size', info_fn=lambda info: update_info)

Update an SMC move-kernel scale from its acceptance diagnostics.

Source code in probjax/inference/smc/tuning.py

def acceptance_rate_adaptor(
    *,
    target_acceptance_rate: float = 0.234,
    scale_key: str = "step_size",
    info_fn=lambda info: info.update_info,
) -> Adaptor:
    """Update an SMC move-kernel scale from its acceptance diagnostics."""
    return _stateless_adaptor(
        lambda params, _state, info: tune_from_kernel_info(
            params,
            info_fn(info),
            target_acceptance_rate=target_acceptance_rate,
            scale_key=scale_key,
        )
    )

probjax.inference.slice_step_size_adaptor ¶

slice_step_size_adaptor(target_num_evaluations, max_evaluations, **kwargs)

Local dual-averaging rule for slice-width adaptation.

Source code in probjax/inference/mcmc/adaptation.py

def slice_step_size_adaptor(
    target_num_evaluations: float,
    max_evaluations: int,
    **kwargs,
) -> Adaptor:
    """Local dual-averaging rule for slice-width adaptation."""
    return step_size_adaptor(
        target=target_num_evaluations / max_evaluations,
        target_from_info_fn=lambda info: info.num_evals / max_evaluations,
        **kwargs,
    )

probjax.inference.particle_adaptor ¶

particle_adaptor(*, updates=None)

Update MCMC geometry from the current SMC particle population.

Source code in probjax/inference/smc/tuning.py

def particle_adaptor(*, updates: Optional[Dict[str, str]] = None) -> Adaptor:
    """Update MCMC geometry from the current SMC particle population."""
    return _stateless_adaptor(
        lambda params, state, _info: tune_from_particles(
            params, state.particles, updates=updates
        )
    )

probjax.inference.adapt `module-attribute` ¶

adapt = adaptor_warmup

probjax.inference.adapt_step ¶

adapt_step(key, kernel, adaptor, state, params, adaptor_state, *args)

Advance a kernel once and apply one local parameter-adaptation update.

Source code in probjax/inference/adaptation.py

@partial(jax.jit, static_argnames=("kernel", "adaptor"))
def adapt_step(
    key, kernel: Kernel, adaptor: Adaptor, state, params, adaptor_state, *args
):
    """Advance a kernel once and apply one local parameter-adaptation update."""
    state, kernel_info = kernel.step(key, state, params, *args)
    adaptor_state, params, adaptor_info = adaptor.update(
        state, kernel_info, adaptor_state, params
    )
    return state, params, adaptor_state, kernel_info, adaptor_info

probjax.inference.as_warmup ¶

as_warmup(adaptor, *, collect=False)

Create the standard fixed-length warmup policy from a local adaptor.

Source code in probjax/inference/adaptation.py

def as_warmup(adaptor: Adaptor, *, collect: bool = False):
    """Create the standard fixed-length warmup policy from a local adaptor."""

    def run(key, kernel, state, params, num_steps, args=None):
        return adaptor_warmup(
            key,
            kernel,
            adaptor,
            state,
            params,
            num_steps,
            args,
            collect=collect,
        )

    return Warmup(run)

Convergence diagnostics¶

probjax.inference.effective_sample_size ¶

effective_sample_size(input_array, chain_axis=0, sample_axis=1)

Compute estimate of the effective sample size (ess).

Parameters¶

input_array: An array representing multiple chains of MCMC samples. The array must contains a chain dimension and a sample dimension. chain_axis The axis indicating the multiple chains. Default to 0. sample_axis The axis indicating a single chain of MCMC samples. Default to 1.

Returns¶

NDArray of the resulting statistics (ess), with the chain and sample dimensions squeezed. Variables whose within-chain variance is numerically zero have an effective sample size of zero. Variables containing a NaN draw are undefined and report NaN; the reduction is per variable, so an independent finite variable is unaffected. Infinities are ordered values and are not treated as missing.

Notes¶

The basic ess (:math:N_{\mathit{eff}}) diagnostic is computed by:

.. math:: \hat{N}_{\mathit{eff}} = \frac{MN}{\hat{\tau}}

.. math:: \hat{\tau} = -1 + 2 \sum_{t'=0}^K \hat{P}_{t'}

where :math:M is the number of chains, :math:N the number of draws, :math:\hat{\rho}_t is the estimated _autocorrelation at lag :math:t, and :math:K is the last integer for which :math:\hat{P}_{K} = \hat{\rho}_{2K} + \hat{\rho}_{2K+1} is still positive :cite:p:stan_ess,gelman1995bayesian.

The current implementation is similar to Stan, which uses Geyer's initial monotone sequence criterion :cite:p:geyer1992practical,geyer2011introduction.

Source code in blackjax/diagnostics.py

def effective_sample_size(
    input_array: ArrayLike, chain_axis: int = 0, sample_axis: int = 1
) -> Array:
    """Compute estimate of the effective sample size (ess).

    Parameters
    ----------
    input_array:
        An array representing multiple chains of MCMC samples. The array must
        contains a chain dimension and a sample dimension.
    chain_axis
        The axis indicating the multiple chains. Default to 0.
    sample_axis
        The axis indicating a single chain of MCMC samples. Default to 1.

    Returns
    -------
    NDArray of the resulting statistics (ess), with the chain and sample dimensions squeezed.
    Variables whose within-chain variance is numerically zero have an effective
    sample size of zero.  Variables containing a ``NaN`` draw are undefined and
    report ``NaN``; the reduction is per variable, so an independent finite
    variable is unaffected.  Infinities are ordered values and are not treated
    as missing.

    Notes
    -----
    The basic ess (:math:`N_{\\mathit{eff}}`) diagnostic is computed by:

    .. math:: \\hat{N}_{\\mathit{eff}} = \\frac{MN}{\\hat{\\tau}}

    .. math:: \\hat{\\tau} = -1 + 2 \\sum_{t'=0}^K \\hat{P}_{t'}

    where :math:`M` is the number of chains, :math:`N` the number of draws,
    :math:`\\hat{\\rho}_t` is the estimated _autocorrelation at lag :math:`t`, and
    :math:`K` is the last integer for which :math:`\\hat{P}_{K} = \\hat{\\rho}_{2K} +
    \\hat{\\rho}_{2K+1}` is still positive :cite:p:`stan_ess,gelman1995bayesian`.

    The current implementation is similar to Stan, which uses Geyer's initial monotone sequence
    criterion :cite:p:`geyer1992practical,geyer2011introduction`.

    """
    input_shape = input_array.shape
    sample_axis = sample_axis if sample_axis >= 0 else len(input_shape) + sample_axis
    num_chains = input_shape[chain_axis]
    num_samples = input_shape[sample_axis]
    assert (
        num_samples > 1
    ), f"The input array must have at least 2 samples, got only {num_samples}."

    first_sample = jnp.take(input_array, jnp.array([0]), axis=sample_axis)
    has_within_chain_variation = jnp.any(
        input_array != first_sample,
        axis=(chain_axis, sample_axis),
        keepdims=True,
    )

    mean_across_chain = input_array.mean(axis=sample_axis, keepdims=True)
    # Compute autocovariance estimates for every lag for the input array using FFT.
    centered_array = input_array - mean_across_chain
    m = next_fast_len(2 * num_samples)
    ifft_ary = jnp.fft.rfft(centered_array, n=m, axis=sample_axis)
    ifft_ary *= jnp.conjugate(ifft_ary)
    autocov_value = jnp.fft.irfft(ifft_ary, n=m, axis=sample_axis)
    autocov_value = (
        jnp.take(autocov_value, jnp.arange(num_samples), axis=sample_axis) / num_samples
    )
    mean_autocov_var = autocov_value.mean(chain_axis, keepdims=True)
    mean_var0 = (
        jnp.take(mean_autocov_var, jnp.array([0]), axis=sample_axis)
        * num_samples
        / (num_samples - 1.0)
    )
    is_numerically_degenerate = jnp.isfinite(mean_var0) & (
        ~has_within_chain_variation | (mean_var0 <= 0.0)
    )
    weighted_var = mean_var0 * (num_samples - 1.0) / num_samples
    weighted_var = jax.lax.cond(
        num_chains > 1,
        lambda mean_across_chain: weighted_var
        + mean_across_chain.var(axis=chain_axis, ddof=1, keepdims=True),
        lambda _: weighted_var,
        operand=mean_across_chain,
    )
    weighted_var = jnp.where(
        is_numerically_degenerate, jnp.ones_like(weighted_var), weighted_var
    )

    # Geyer's initial positive sequence
    num_samples_even = num_samples - num_samples % 2
    mean_autocov_var_tp1 = jnp.take(
        mean_autocov_var, jnp.arange(1, num_samples_even), axis=sample_axis
    )
    rho_hat = jnp.concatenate(
        [
            jnp.ones_like(mean_var0),
            1.0 - (mean_var0 - mean_autocov_var_tp1) / weighted_var,
        ],
        axis=sample_axis,
    )

    rho_hat = jnp.moveaxis(rho_hat, sample_axis, 0)
    rho_hat_even = rho_hat[0::2]
    rho_hat_odd = rho_hat[1::2]

    mask0 = (rho_hat_even + rho_hat_odd) > 0.0
    carry_cond = jnp.ones_like(mask0[0])
    max_t = jnp.zeros_like(mask0[0], dtype=int)

    def positive_sequence_body_fn(state, mask_t):
        t, carry_cond, max_t = state
        next_mask = carry_cond & mask_t
        next_max_t = jnp.where(next_mask, jnp.ones_like(max_t) * t, max_t)
        return (t + 1, next_mask, next_max_t), next_mask

    (*_, max_t_next), mask = jax.lax.scan(
        positive_sequence_body_fn, (0, carry_cond, max_t), mask0
    )
    indices = jnp.indices(max_t_next.shape)
    indices = tuple([max_t_next + 1] + [indices[i] for i in range(max_t_next.ndim)])
    rho_hat_odd = jnp.where(mask, rho_hat_odd, jnp.zeros_like(rho_hat_odd))
    # improve estimation
    mask_even = mask.at[indices].set(rho_hat_even[indices] > 0)
    rho_hat_even = jnp.where(mask_even, rho_hat_even, jnp.zeros_like(rho_hat_even))

    # Geyer's initial monotone sequence
    def monotone_sequence_body_fn(rho_hat_sum_tm1, rho_hat_sum_t):
        update_mask = rho_hat_sum_t > rho_hat_sum_tm1
        next_rho_hat_sum_t = jnp.where(update_mask, rho_hat_sum_tm1, rho_hat_sum_t)
        return next_rho_hat_sum_t, (update_mask, next_rho_hat_sum_t)

    rho_hat_sum = rho_hat_even + rho_hat_odd
    _, (update_mask, update_value) = jax.lax.scan(
        monotone_sequence_body_fn, rho_hat_sum[0], rho_hat_sum
    )

    rho_hat_even_final = jnp.where(update_mask, update_value / 2.0, rho_hat_even)
    rho_hat_odd_final = jnp.where(update_mask, update_value / 2.0, rho_hat_odd)

    # compute effective sample size
    ess_raw = num_chains * num_samples
    tau_hat = (
        -1.0
        + 2.0 * jnp.sum(rho_hat_even_final + rho_hat_odd_final, axis=0)
        - rho_hat_even_final[indices]
    )

    tau_hat = jnp.maximum(tau_hat, 1 / np.log10(ess_raw))
    ess = ess_raw / tau_hat
    ess = jnp.where(is_numerically_degenerate.squeeze(), 0.0, ess.squeeze())

    # A NaN draw is a missing observation, not a value the estimator can use.
    # Geyer's truncation above gates on partial sums being > 0, and every such
    # comparison is False for NaN, so the truncation collapses and tau_hat
    # falls back to its 1 / log10(MN) floor — manufacturing a large, entirely
    # fictitious sample size.  Report the component as undefined instead.  The
    # reduction runs over the chain and draw axes only, so a contaminated
    # component never poisons an independent finite one, and ``.squeeze()``
    # mirrors the squeeze applied to ``ess`` just above.
    has_nan = jnp.any(jnp.isnan(input_array), axis=(chain_axis, sample_axis))
    ess = jnp.where(jnp.squeeze(has_nan), jnp.nan, ess)

    return ess

probjax.inference.ess_bulk ¶

ess_bulk(input_array, chain_axis=0, sample_axis=1)

Bulk effective sample size (rank-normalized split-chain ESS).

Computes the bulk ESS from Vehtari et al. (2021): rank-normalizes draws after splitting each chain in half, then applies the standard autocorrelation-based :func:effective_sample_size estimator. This diagnostic is robust to non-stationarity and multimodality.

Parameters¶

input_array An array representing multiple chains of MCMC samples. The array must contain a chain dimension and a sample dimension. chain_axis The axis indicating the multiple chains. Default 0. sample_axis The axis indicating a single chain of MCMC samples. Default 1.

Returns¶

NDArray of the resulting bulk-ESS, with chain and sample dimensions squeezed.

Notes¶

Algorithm:

Split each chain in half → 2× chains.
Pool all draws and rank-normalize with :math:z_r = \Phi^{-1}((r-3/8)/(n+1/4)).
Apply :func:effective_sample_size to the rank-normalized draws.

References¶

.. cite:p:vehtari2021rank

Source code in blackjax/diagnostics.py

def ess_bulk(
    input_array: ArrayLike, chain_axis: int = 0, sample_axis: int = 1
) -> Array:
    """Bulk effective sample size (rank-normalized split-chain ESS).

    Computes the bulk ESS from Vehtari et al. (2021): rank-normalizes draws
    after splitting each chain in half, then applies the standard
    autocorrelation-based :func:`effective_sample_size` estimator.  This
    diagnostic is robust to non-stationarity and multimodality.

    Parameters
    ----------
    input_array
        An array representing multiple chains of MCMC samples. The array must
        contain a chain dimension and a sample dimension.
    chain_axis
        The axis indicating the multiple chains. Default 0.
    sample_axis
        The axis indicating a single chain of MCMC samples. Default 1.

    Returns
    -------
    NDArray of the resulting bulk-ESS, with chain and sample dimensions squeezed.

    Notes
    -----
    Algorithm:

    1. Split each chain in half → 2× chains.
    2. Pool all draws and rank-normalize with :math:`z_r = \\Phi^{-1}((r-3/8)/(n+1/4))`.
    3. Apply :func:`effective_sample_size` to the rank-normalized draws.

    References
    ----------
    .. cite:p:`vehtari2021rank`
    """
    x = _to_standard_axes(jnp.asarray(input_array), chain_axis, sample_axis)
    x_split = _split_chains(x)
    x_rn = _rank_normalize(x_split)
    # A contaminated component is already all-NaN after _rank_normalize, and
    # effective_sample_size propagates that on its own.
    return effective_sample_size(x_rn)

probjax.inference.ess_tail ¶

ess_tail(input_array, chain_axis=0, sample_axis=1, prob=0.9)

Tail effective sample size.

Computes the tail ESS from Vehtari et al. (2021) as the minimum of the ESS of the lower- and upper-tail indicator functions applied to split-chain draws.

The tail quantiles are determined by prob: the lower tail uses the (1 - prob) / 2 quantile and the upper tail uses the (1 + prob) / 2 quantile. The default prob=0.90 corresponds to the 5^th/95^th percentiles, matching ArviZ's default prob=(0.05, 0.95).

.. warning::

The agreement with az.ess(method="tail") holds for continuous draws only. The upper-tail indicator here is :math:\mathbf{1}(x \ge q_{\text{high}}), whereas Vehtari et al. and ArviZ use :math:\mathbf{1}(x \le q) for both tails. On continuous draws the two are exact complements and the ESS is identical, but on tied draws they are not: for iid Bernoulli(0.1), :math:P(x \ge q_{95}) is 0.097 rather than 0.05, and for a 5-level grid it is 0.21. This is a separate known defect in the tail estimator's tie handling, tracked independently of the rank-normalization fix; note that simply switching to <= does not resolve it, since that indicator is identically 1 on such draws and would be reported as degenerate.

Parameters¶

input_array An array representing multiple chains of MCMC samples. The array must contain a chain dimension and a sample dimension. chain_axis The axis indicating the multiple chains. Default 0. sample_axis The axis indicating a single chain of MCMC samples. Default 1. prob Central-interval probability that determines the tail quantiles. Lower quantile: (1 - prob) / 2; upper quantile: (1 + prob) / 2. Default 0.90 gives the 5^th/95^th-percentile tail, matching az.ess(method="tail") (ArviZ default).

Returns¶

NDArray of the resulting tail-ESS, with chain and sample dimensions squeezed.

Notes¶

Algorithm:

Split each chain in half → 2× chains.
Compute pooled lower/upper quantiles (at (1-prob)/2 and (1+prob)/2) across all split chains and draws.
Form indicator series :math:\mathbf{1}(x \le q_{\text{low}}) and :math:\mathbf{1}(x \ge q_{\text{high}}).
Compute :func:effective_sample_size for each indicator.
Return :math:\min(\text{ESS}_\text{lower}, \text{ESS}_\text{upper}).

References¶

.. cite:p:vehtari2021rank

Source code in blackjax/diagnostics.py

def ess_tail(
    input_array: ArrayLike,
    chain_axis: int = 0,
    sample_axis: int = 1,
    prob: float = 0.90,
) -> Array:
    """Tail effective sample size.

    Computes the tail ESS from Vehtari et al. (2021) as the minimum of the
    ESS of the lower- and upper-tail indicator functions applied to split-chain
    draws.

    The tail quantiles are determined by ``prob``: the lower tail uses the
    ``(1 - prob) / 2`` quantile and the upper tail uses the
    ``(1 + prob) / 2`` quantile.  The default ``prob=0.90`` corresponds to
    the 5th/95th percentiles, matching ArviZ's default ``prob=(0.05, 0.95)``.

    .. warning::

       The agreement with ``az.ess(method="tail")`` holds for **continuous**
       draws only.  The upper-tail indicator here is
       :math:`\\mathbf{1}(x \\ge q_{\\text{high}})`, whereas Vehtari et al.
       and ArviZ use :math:`\\mathbf{1}(x \\le q)` for both tails.  On
       continuous draws the two are exact complements and the ESS is
       identical, but on tied draws they are not: for iid Bernoulli(0.1),
       :math:`P(x \\ge q_{95})` is 0.097 rather than 0.05, and for a 5-level
       grid it is 0.21.  This is a separate known defect in the tail
       estimator's tie handling, tracked independently of the
       rank-normalization fix; note that simply switching to ``<=`` does not
       resolve it, since that indicator is identically 1 on such draws and
       would be reported as degenerate.

    Parameters
    ----------
    input_array
        An array representing multiple chains of MCMC samples. The array must
        contain a chain dimension and a sample dimension.
    chain_axis
        The axis indicating the multiple chains. Default 0.
    sample_axis
        The axis indicating a single chain of MCMC samples. Default 1.
    prob
        Central-interval probability that determines the tail quantiles.
        Lower quantile: ``(1 - prob) / 2``; upper quantile:
        ``(1 + prob) / 2``.  Default ``0.90`` gives the 5th/95th-percentile
        tail, matching ``az.ess(method="tail")`` (ArviZ default).

    Returns
    -------
    NDArray of the resulting tail-ESS, with chain and sample dimensions squeezed.

    Notes
    -----
    Algorithm:

    1. Split each chain in half → 2× chains.
    2. Compute pooled lower/upper quantiles (at ``(1-prob)/2`` and
       ``(1+prob)/2``) across all split chains and draws.
    3. Form indicator series :math:`\\mathbf{1}(x \\le q_{\\text{low}})` and
       :math:`\\mathbf{1}(x \\ge q_{\\text{high}})`.
    4. Compute :func:`effective_sample_size` for each indicator.
    5. Return :math:`\\min(\\text{ESS}_\\text{lower}, \\text{ESS}_\\text{upper})`.

    References
    ----------
    .. cite:p:`vehtari2021rank`
    """
    x = _to_standard_axes(jnp.asarray(input_array), chain_axis, sample_axis)
    x_split = _split_chains(x)

    nchains, nsamples = x_split.shape[0], x_split.shape[1]
    extra_shape = x_split.shape[2:]

    # Tail quantiles derived from the central-interval probability.
    q_low = (1.0 - prob) / 2.0
    q_high = (1.0 + prob) / 2.0

    # Pooled quantiles over all chains and draws, per trailing dimension.
    x_flat = x_split.reshape(nchains * nsamples, *extra_shape)
    q_lo = jnp.quantile(x_flat, q_low, axis=0)
    q_hi = jnp.quantile(x_flat, q_high, axis=0)

    # Indicator series (float for ESS computation).
    # Broadcast q_lo/q_hi over the (nchains, nsamples) leading axes.
    I_lower = (x_split <= q_lo[None, None]).astype(float)
    I_upper = (x_split >= q_hi[None, None]).astype(float)

    ess_lower = effective_sample_size(I_lower)
    ess_upper = effective_sample_size(I_upper)

    # A NaN draw makes the pooled quantiles NaN, every indicator comparison
    # False, and the resulting all-zero series degenerate — which would report
    # 0 rather than "undefined".  Guard once on the reduced value: the
    # indicators themselves are never NaN, so the estimator cannot propagate
    # this on its own the way it does for ess_bulk.
    return _propagate_nan_components(jnp.minimum(ess_lower, ess_upper), x_split)

probjax.inference.potential_scale_reduction ¶

potential_scale_reduction(input_array, chain_axis=0, sample_axis=1)

Gelman and Rubin (1992)'s potential scale reduction for computing multiple MCMC chain convergence.

Parameters¶

input_array: An array representing multiple chains of MCMC samples. The array must contains a chain dimension and a sample dimension. chain_axis The axis indicating the multiple chains. Default to 0. sample_axis The axis indicating a single chain of MCMC samples. Default to 1.

Returns¶

NDArray of the resulting statistics (r-hat), with the chain and sample dimensions squeezed.

Notes¶

The diagnostic is computed by:

.. math:: \hat{R} = \sqrt{\frac{\hat{V}}{W}}

where :math:W is the within-chain variance and :math:\hat{V} is the posterior variance estimate for the pooled traces. This is the potential scale reduction factor, which converges to unity when each of the traces is a sample from the target posterior. Values greater than one indicate that one or more chains have not yet converged :cite:p:stan_rhat,gelman1992inference.

Source code in blackjax/diagnostics.py

def potential_scale_reduction(
    input_array: ArrayLike, chain_axis: int = 0, sample_axis: int = 1
) -> Array:
    """Gelman and Rubin (1992)'s potential scale reduction for computing multiple MCMC chain convergence.

    Parameters
    ----------
    input_array:
        An array representing multiple chains of MCMC samples. The array must
        contains a chain dimension and a sample dimension.
    chain_axis
        The axis indicating the multiple chains. Default to 0.
    sample_axis
        The axis indicating a single chain of MCMC samples. Default to 1.

    Returns
    -------
    NDArray of the resulting statistics (r-hat), with the chain and sample dimensions squeezed.

    Notes
    -----
    The diagnostic is computed by:

    .. math:: \\hat{R} = \\sqrt{\\frac{\\hat{V}}{W}}

    where :math:`W` is the within-chain variance and :math:`\\hat{V}` is the posterior variance
    estimate for the pooled traces. This is the potential scale reduction factor, which
    converges to unity when each of the traces is a sample from the target posterior. Values
    greater than one indicate that one or more chains have not yet converged :cite:p:`stan_rhat,gelman1992inference`.

    """
    assert (
        input_array.shape[chain_axis] > 1
    ), "potential_scale_reduction as implemented only works for two or more chains."

    num_samples = input_array.shape[sample_axis]
    # Compute stats for each chain
    per_chain_mean = input_array.mean(axis=sample_axis, keepdims=True)
    per_chain_var = input_array.var(axis=sample_axis, ddof=1, keepdims=True)
    # Compute between-chain stats
    between_chain_variance = num_samples * per_chain_mean.var(
        axis=chain_axis, ddof=1, keepdims=True
    )
    # Compute within-chain stats
    within_chain_variance = per_chain_var.mean(axis=chain_axis, keepdims=True)
    # Estimate of marginal posterior variance
    rhat_value = jnp.sqrt(
        (between_chain_variance / within_chain_variance + num_samples - 1)
        / (num_samples)
    )
    return rhat_value.squeeze()

probjax.inference.rhat ¶

rhat(input_array, chain_axis=0, sample_axis=1)

Rank-normalized split-R̂ (Vehtari et al. 2021).

The modern improved R̂ diagnostic. Combines two split-chain R̂ values — one on rank-normalized draws and one on rank-normalized folded draws — and returns the maximum. The folded component catches scale/variance non-convergence that the bulk component can miss.

This matches the default az.rhat(method="rank") convention in ArviZ.

Parameters¶

input_array An array representing multiple chains of MCMC samples. The array must contain a chain dimension and a sample dimension. At least 2 chains and at least 4 draws per chain are required. chain_axis The axis indicating the multiple chains. Default 0. sample_axis The axis indicating a single chain of MCMC samples. Default 1.

Returns¶

NDArray of the resulting R̂ values, with chain and sample dimensions squeezed. Values close to 1.0 indicate convergence; values above 1.01 suggest chains have not converged.

Notes¶

Algorithm (Vehtari et al. 2021, § 4):

Split each chain in half → 2× chains.
Rank-normalize with the Blom plotting position :math:z_r = \Phi^{-1}((r - 3/8) / (n + 1/4)) over the joint pool.
Compute the standard split-R̂ on the rank-normalized draws (bulk).
Compute the folded draws :math:|x - \mathrm{median}(x)|, rank-normalize them, and compute split-R̂ again (tail).
Return :math:\max(\hat{R}_{\text{bulk}}, \hat{R}_{\text{tail}}).

.. warning::

NaN from this function does not uniquely mean "the draws contained a missing observation". The folded component subtracts the pooled median, so when at least half of a component's pooled draws are +inf (or -inf) the median is infinite and the fold evaluates inf - inf, which is NaN. Such a component returns NaN from :func:rhat while containing no NaN at all — and while :func:ess_bulk and :func:ess_tail still return finite values for it. Infinity handling is a separate open question from the missing-data contract; this is documented, not yet decided.

References¶

.. cite:p:vehtari2021rank

Source code in blackjax/diagnostics.py

def rhat(input_array: ArrayLike, chain_axis: int = 0, sample_axis: int = 1) -> Array:
    """Rank-normalized split-R̂ (Vehtari et al. 2021).

    The modern improved R̂ diagnostic.  Combines two split-chain R̂ values —
    one on rank-normalized draws and one on rank-normalized *folded* draws —
    and returns the maximum.  The folded component catches scale/variance
    non-convergence that the bulk component can miss.

    This matches the default ``az.rhat(method="rank")`` convention in ArviZ.

    Parameters
    ----------
    input_array
        An array representing multiple chains of MCMC samples. The array must
        contain a chain dimension and a sample dimension.  At least 2 chains
        and at least 4 draws per chain are required.
    chain_axis
        The axis indicating the multiple chains. Default 0.
    sample_axis
        The axis indicating a single chain of MCMC samples. Default 1.

    Returns
    -------
    NDArray of the resulting R̂ values, with chain and sample dimensions
    squeezed.  Values close to 1.0 indicate convergence; values above 1.01
    suggest chains have not converged.

    Notes
    -----
    Algorithm (Vehtari et al. 2021, § 4):

    1. Split each chain in half → 2× chains.
    2. Rank-normalize with the Blom plotting position
       :math:`z_r = \\Phi^{-1}((r - 3/8) / (n + 1/4))` over the joint pool.
    3. Compute the standard split-R̂ on the rank-normalized draws (**bulk**).
    4. Compute the folded draws :math:`|x - \\mathrm{median}(x)|`, rank-normalize
       them, and compute split-R̂ again (**tail**).
    5. Return :math:`\\max(\\hat{R}_{\\text{bulk}}, \\hat{R}_{\\text{tail}})`.

    .. warning::

       ``NaN`` from this function does **not** uniquely mean "the draws
       contained a missing observation".  The folded component subtracts the
       pooled median, so when at least half of a component's pooled draws are
       ``+inf`` (or ``-inf``) the median is infinite and the fold evaluates
       ``inf - inf``, which is ``NaN``.  Such a component returns ``NaN`` from
       :func:`rhat` while containing no ``NaN`` at all — and while
       :func:`ess_bulk` and :func:`ess_tail` still return finite values for it.
       Infinity handling is a separate open question from the missing-data
       contract; this is documented, not yet decided.

    References
    ----------
    .. cite:p:`vehtari2021rank`
    """
    x = _to_standard_axes(jnp.asarray(input_array), chain_axis, sample_axis)
    # Split each chain in half → (2*nchains, nsamples//2, …).
    x_split = _split_chains(x)

    # Bulk: rank-normalize split draws, then compute split-R̂.
    x_rn = _rank_normalize(x_split)
    rhat_bulk = potential_scale_reduction(x_rn, chain_axis=0, sample_axis=1)

    # Tail: fold the split draws about their joint median, rank-normalize,
    # then compute split-R̂.  Catches variance non-convergence.
    nchains_split = x_split.shape[0]
    nsamples_split = x_split.shape[1]
    extra_shape = x_split.shape[2:]
    x_flat = x_split.reshape(nchains_split * nsamples_split, *extra_shape)
    # Global median per trailing dimension (scalar when extra_shape is empty).
    x_folded = jnp.abs(x_split - jnp.median(x_flat, axis=0))
    x_folded_rn = _rank_normalize(x_folded)
    rhat_tail = potential_scale_reduction(x_folded_rn, chain_axis=0, sample_axis=1)

    return jnp.maximum(rhat_bulk, rhat_tail)

probjax.inference.pareto_khat ¶

pareto_khat(x, tail='both', tail_frac=0.1)

Pareto shape parameter k̂ for tail diagnosis.

Fits a Generalised Pareto Distribution (GPD) to the upper and/or lower tail of a 1-D sample and returns the estimated shape parameter k̂.

Parameters¶

x 1-D array of draws (or any array; it is ravelled before use). tail Which tail to fit: "upper", "lower", or "both" (default). When "both", returns the maximum of the two k̂ estimates. tail_frac Fraction of samples used as the tail. Default 0.10 (10 %). A minimum of 5 tail samples is always enforced.

Returns¶

Scalar Array: the Pareto shape estimate k̂. Values below 0.5 indicate reliable tail estimates; 0.5–0.7 are moderate; above 0.7 may be unreliable.

Notes¶

Uses the Zhang & Stephens (2009) empirical-Bayes estimator implemented in the internal :func:_gpdfit. The upper tail is modelled directly; the lower tail is reflected and modelled as an upper tail.

Source code in blackjax/diagnostics.py

def pareto_khat(x: ArrayLike, tail: str = "both", tail_frac: float = 0.10) -> Array:
    """Pareto shape parameter k̂ for tail diagnosis.

    Fits a Generalised Pareto Distribution (GPD) to the upper and/or lower
    tail of a 1-D sample and returns the estimated shape parameter k̂.

    Parameters
    ----------
    x
        1-D array of draws (or any array; it is ravelled before use).
    tail
        Which tail to fit: ``"upper"``, ``"lower"``, or ``"both"`` (default).
        When ``"both"``, returns the maximum of the two k̂ estimates.
    tail_frac
        Fraction of samples used as the tail. Default 0.10 (10 %).
        A minimum of 5 tail samples is always enforced.

    Returns
    -------
    Scalar Array: the Pareto shape estimate k̂.  Values below 0.5 indicate
    reliable tail estimates; 0.5–0.7 are moderate; above 0.7 may be
    unreliable.

    Notes
    -----
    Uses the Zhang & Stephens (2009) empirical-Bayes estimator implemented
    in the internal :func:`_gpdfit`.  The upper tail is modelled directly;
    the lower tail is reflected and modelled as an upper tail.
    """
    x_flat = jnp.asarray(x).ravel()
    n = x_flat.shape[0]
    tail_size = max(int(tail_frac * n), 5)

    x_sorted = jnp.sort(x_flat)  # ascending

    if tail in ("upper", "both"):
        upper_tail = x_sorted[n - tail_size :]
        threshold_upper = x_sorted[n - tail_size - 1]
        exc_upper = upper_tail - threshold_upper  # >= 0, ascending
        k_upper, _ = _gpdfit(exc_upper)

    if tail in ("lower", "both"):
        # Reflect the lower tail so it becomes an upper-tail problem.
        lower_tail_reflected = -x_sorted[:tail_size][::-1]  # ascending
        threshold_lower = -x_sorted[tail_size]
        exc_lower = lower_tail_reflected - threshold_lower  # >= 0, ascending
        k_lower, _ = _gpdfit(exc_lower)

    if tail == "upper":
        return k_upper
    if tail == "lower":
        return k_lower
    return jnp.maximum(k_upper, k_lower)

probjax.inference.kernelized_stein_discrepancy ¶

kernelized_stein_discrepancy(samples, grad_logdensity_fn, kernel, *, statistic='v')

Estimate squared kernelized Stein discrepancy with a V- or U-statistic.

Parameters¶

samples PyTree (an array, or a dict/tuple/NamedTuple of per-field arrays) whose leaves share a leading num_samples axis. Chain dimensions must be pooled explicitly before calling this function. grad_logdensity_fn Target score function, mapping one sample (one leading-axis slice of samples, with the same PyTree structure) to the gradient of its log density. For example, jax.grad(logdensity_fn). The normalizing constant is not required. kernel Scalar-valued, twice differentiable positive semidefinite kernel taking two flattened sample vectors. Bind kernel parameters such as bandwidth with functools.partial. The kernel must satisfy the target's Stein boundary conditions; choosing an appropriate kernel is the caller's responsibility.

statistic "v" includes diagonal pairs; "u" excludes them and requires at least two draws. The U-statistic can be negative, even in exact arithmetic.

Returns¶

Estimate of squared KSD, not a calibrated p-value or convergence certificate. The V-statistic has a sample-dependent diagonal bias of order 1/n.

Notes¶

Complements rather than replaces R-hat. Approximately independent draws are assumed; autocorrelation and an unvisited mode can make the result misleading. A positive semidefinite kernel alone does not identify every distribution: this requires additional conditions, including integral strict positive definiteness :cite:p:liu2016kernelized. Gaussian RBF kernels do not in general detect non-convergence in dimension three or higher. IMQ kernels have convergence guarantees for suitable targets, but targets with bounded scores, including some heavy-tailed targets, remain a gap :cite:p:gorham2017kernels.

Uses the score/kernel derivative expression in :cite:p:liu2016kernelized. Rows are processed sequentially and columns are vectorized. The computation has quadratic cost in sample count and mixed Hessians with d² entries; temporary pairwise storage scales as O(n*d²) rather than O(n²*d²). The supplied functions must support JAX differentiation and transformations. No clipping is applied to either statistic.

Vendored from unreleased BlackJAX main; see the module docstring.

Source code in probjax/inference/diagnostics.py

def kernelized_stein_discrepancy(
    samples: ArrayLikeTree,
    grad_logdensity_fn: Callable,
    kernel: Callable,
    *,
    statistic: str = "v",
) -> Array:
    """Estimate squared kernelized Stein discrepancy with a V- or U-statistic.

    Parameters
    ----------
    samples
        PyTree (an array, or a dict/tuple/NamedTuple of per-field arrays) whose
        leaves share a leading ``num_samples`` axis. Chain dimensions must be
        pooled explicitly before calling this function.
    grad_logdensity_fn
        Target score function, mapping one sample (one leading-axis slice of
        ``samples``, with the same PyTree structure) to the gradient of its log
        density. For example, ``jax.grad(logdensity_fn)``. The normalizing
        constant is not required.
    kernel
        Scalar-valued, twice differentiable positive semidefinite kernel taking
        two flattened sample vectors. Bind kernel parameters such as bandwidth
        with ``functools.partial``. The kernel must satisfy the target's Stein
        boundary conditions; choosing an appropriate kernel is the caller's
        responsibility.

    statistic
        ``"v"`` includes diagonal pairs; ``"u"`` excludes them and requires at
        least two draws. The U-statistic can be negative, even in exact arithmetic.

    Returns
    -------
    Estimate of squared KSD, not a calibrated p-value or convergence certificate.
    The V-statistic has a sample-dependent diagonal bias of order ``1/n``.

    Notes
    -----
    Complements rather than replaces R-hat. Approximately independent draws are
    assumed; autocorrelation and an unvisited mode can make the result misleading.
    A positive semidefinite kernel alone does not identify every distribution:
    this requires additional conditions, including integral strict positive
    definiteness :cite:p:`liu2016kernelized`. Gaussian RBF kernels do not in
    general detect non-convergence in dimension three or higher. IMQ kernels
    have convergence guarantees for suitable targets, but targets with bounded
    scores, including some heavy-tailed targets, remain a gap
    :cite:p:`gorham2017kernels`.

    Uses the score/kernel derivative expression in :cite:p:`liu2016kernelized`.
    Rows are processed sequentially and columns are vectorized. The computation
    has quadratic cost in sample count and mixed Hessians with ``d²`` entries;
    temporary pairwise storage scales as ``O(n*d²)`` rather than ``O(n²*d²)``.
    The supplied functions must support JAX differentiation and transformations.
    No clipping is applied to either statistic.

    Vendored from unreleased BlackJAX main; see the module docstring.
    """
    if statistic not in ("v", "u"):
        raise ValueError("statistic must be v or u")
    flat = jax.vmap(lambda s: ravel_pytree(s)[0])(samples)
    scores = jax.vmap(lambda s: ravel_pytree(grad_logdensity_fn(s))[0])(samples)
    kernel_value_and_grad = jax.value_and_grad(kernel, argnums=(0, 1))
    mixed_derivative = jax.jacfwd(jax.grad(kernel, argnums=0), argnums=1)

    def stein_kernel(x, y, score_x, score_y):
        value, (grad_x, grad_y) = kernel_value_and_grad(x, y)
        return (
            value * jnp.dot(score_x, score_y)
            + jnp.dot(score_x, grad_y)
            + jnp.dot(grad_x, score_y)
            + jnp.trace(mixed_derivative(x, y))
        )

    def row_sum(args):
        x, score_x = args
        return jax.vmap(stein_kernel, in_axes=(None, 0, None, 0))(
            x, flat, score_x, scores
        ).sum()

    n = flat.shape[0]
    total = jax.lax.map(row_sum, (flat, scores)).sum()
    if statistic == "v":
        return total / n**2
    diagonal = jax.vmap(stein_kernel)(flat, flat, scores, scores).sum()
    return (total - diagonal) / (n * (n - 1))

probjax.inference.imq_kernel ¶

imq_kernel(x, y, c=1.0, beta=-0.5)

Inverse multiquadric kernel (c² + ||x-y||²)**beta.

Use c > 0 and -1 < beta < 0. Bind non-default parameters with functools.partial. Convergence guarantees additionally require target conditions; see :cite:p:gorham2017kernels.

Vendored from unreleased BlackJAX main; see the module docstring.

Source code in probjax/inference/diagnostics.py

def imq_kernel(x: Array, y: Array, c: float = 1.0, beta: float = -0.5) -> Array:
    """Inverse multiquadric kernel ``(c² + ||x-y||²)**beta``.

    Use ``c > 0`` and ``-1 < beta < 0``. Bind non-default parameters with
    ``functools.partial``. Convergence guarantees additionally require target
    conditions; see :cite:p:`gorham2017kernels`.

    Vendored from unreleased BlackJAX main; see the module docstring.
    """
    return (c**2 + jnp.sum((x - y) ** 2)) ** beta

probjax.inference.divergence_concentration ¶

divergence_concentration(is_divergent, *, rate_threshold=0.02)

Flag a run where sampling-phase divergences concentrate on a minority of chains.

Parameters¶

is_divergent Per-draw divergence flags (bool or 0/1) for the sampling (post-warmup) phase only, shape (n_chains, n_draws). rate_threshold Minimum per-chain divergence rate for chain k to count as flagged. Default 0.02 (2%).

Returns¶

:class:DivergenceConcentrationReport. warn is true iff between 1 and max(1, n_chains // 4) chains are flagged -- a minority-outlier trigger; an ensemble where most/all chains cross the threshold returns populated fields but no warning.

Notes¶

Counting is sampling-phase only; warmup divergences are out of scope.
Concatenating warmup draws in front of sampling draws dilutes a real signal below threshold rather than raising a false alarm -- slice to sampling draws only before calling this.
multinomial_p_value is context only; it never decides warn.

Source code in blackjax/diagnostics.py

def divergence_concentration(
    is_divergent: ArrayLike,
    *,
    rate_threshold: float = 0.02,
) -> DivergenceConcentrationReport:
    """Flag a run where sampling-phase divergences concentrate on a minority of chains.

    Parameters
    ----------
    is_divergent
        Per-draw divergence flags (bool or 0/1) for the sampling
        (post-warmup) phase only, shape ``(n_chains, n_draws)``.
    rate_threshold
        Minimum per-chain divergence rate for chain ``k`` to count as
        flagged. Default ``0.02`` (2%).

    Returns
    -------
    :class:`DivergenceConcentrationReport`. ``warn`` is true iff between
    1 and ``max(1, n_chains // 4)`` chains are flagged -- a
    minority-outlier trigger; an ensemble where most/all chains cross the
    threshold returns populated fields but no warning.

    Notes
    -----
    - Counting is sampling-phase only; warmup divergences are out of scope.
    - Concatenating warmup draws in front of sampling draws dilutes a real
      signal below threshold rather than raising a false alarm -- slice to
      sampling draws only before calling this.
    - ``multinomial_p_value`` is context only; it never decides ``warn``.
    """
    # NaN must flag, not silently pass.
    is_divergent = jnp.asarray(is_divergent).astype(bool)
    n_draws = is_divergent.shape[-1]
    num_chains = is_divergent.shape[-2]
    counts = jnp.sum(is_divergent, axis=-1)

    quarter = max(n_draws // 4, 1)
    early_rate = jnp.sum(is_divergent[:, :quarter], axis=-1) / quarter
    late_rate = jnp.sum(is_divergent[:, n_draws - quarter :], axis=-1) / quarter

    return _divergence_concentration_core(
        counts,
        n_draws,
        num_chains,
        early_rate,
        late_rate,
        rate_threshold=rate_threshold,
    )

probjax.inference.divergence_concentration_from_counts ¶

divergence_concentration_from_counts(chain_divergence_counts, n_draws, *, rate_threshold=0.02)

Same as :func:divergence_concentration, from precomputed per-chain counts.

Use when per-draw flags are not retained but per-chain totals are. early_rate / late_rate are NaN (no per-draw resolution for a quarter profile); :func:format_divergence_warning omits that part of the message rather than printing "nan%".

Parameters¶

chain_divergence_counts Per-chain divergence counts for the sampling phase, shape (n_chains,). n_draws Number of sampling draws per chain. rate_threshold See :func:divergence_concentration.

Returns¶

:class:DivergenceConcentrationReport

Source code in blackjax/diagnostics.py

def divergence_concentration_from_counts(
    chain_divergence_counts: ArrayLike,
    n_draws: int,
    *,
    rate_threshold: float = 0.02,
) -> DivergenceConcentrationReport:
    """Same as :func:`divergence_concentration`, from precomputed per-chain counts.

    Use when per-draw flags are not retained but per-chain totals are.
    ``early_rate`` / ``late_rate`` are ``NaN`` (no per-draw resolution for
    a quarter profile); :func:`format_divergence_warning` omits that part
    of the message rather than printing "nan%".

    Parameters
    ----------
    chain_divergence_counts
        Per-chain divergence counts for the sampling phase, shape
        ``(n_chains,)``.
    n_draws
        Number of sampling draws per chain.
    rate_threshold
        See :func:`divergence_concentration`.

    Returns
    -------
    :class:`DivergenceConcentrationReport`
    """
    # NaN must flag (as fully divergent), not silently pass.
    counts = jnp.asarray(chain_divergence_counts).astype(float)
    counts = jnp.clip(jnp.nan_to_num(counts, nan=float(n_draws)), 0.0, float(n_draws))
    num_chains = counts.shape[-1]
    nan_quarters = jnp.full((num_chains,), jnp.nan)
    return _divergence_concentration_core(
        counts,
        n_draws,
        num_chains,
        nan_quarters,
        nan_quarters,
        rate_threshold=rate_threshold,
    )

probjax.inference.format_divergence_warning ¶

format_divergence_warning(report)

Render a :class:DivergenceConcentrationReport as a human-readable message.

Returns "" when report.warn is false. Otherwise returns one sentence per flagged chain (newline-separated when more than one chain is flagged). Not JIT-compatible by design — call it on concrete report values (e.g. after a sampling run has completed), not inside traced code.

Parameters¶

report A report produced by :func:divergence_concentration or :func:divergence_concentration_from_counts.

Returns¶

str, empty when there is nothing to warn about.

Source code in blackjax/diagnostics.py

def format_divergence_warning(report: DivergenceConcentrationReport) -> str:
    """Render a :class:`DivergenceConcentrationReport` as a human-readable message.

    Returns ``""`` when ``report.warn`` is false. Otherwise returns one
    sentence per flagged chain (newline-separated when more than one
    chain is flagged). Not JIT-compatible by design — call it on concrete
    report values (e.g. after a sampling run has completed), not inside
    traced code.

    Parameters
    ----------
    report
        A report produced by :func:`divergence_concentration` or
        :func:`divergence_concentration_from_counts`.

    Returns
    -------
    ``str``, empty when there is nothing to warn about.
    """
    if not bool(report.warn):
        return ""

    lines = []
    flagged = np.asarray(report.flagged)
    for k in np.flatnonzero(flagged):
        k = int(k)
        pct = "%.1f" % (float(report.rates[k]) * 100.0)
        median = "%.1f" % (float(report.median_other_rate[k]) * 100.0)
        early = float(report.early_rate[k])
        late = float(report.late_rate[k])
        if np.isfinite(early) and np.isfinite(late):
            early_str = "%.1f" % (early * 100.0)
            late_str = "%.1f" % (late * 100.0)
            quarters = " ({}% in the first quarter, {}% in the last)".format(
                early_str,
                late_str,
            )
        else:
            quarters = ""
        lines.append(
            f"chain {k}: {pct}% of its sampling draws diverged{quarters} — "
            f"significantly more than the other chains ({median}% median)"
            "; draws from this chain, especially early ones, may be "
            "untrustworthy."
        )
    return "\n".join(lines)

probjax.inference.DivergenceConcentrationReport ¶

Bases: NamedTuple

Per-run divergence-concentration diagnostic report.

All fields are plain JAX numerics — no strings — so the statistic itself stays JIT-compilable. Pass the report to :func:format_divergence_warning to render a human-readable message from concrete (non-traced) values.

Attributes¶

warn Bool. Whether a minority of chains (see the module notes on :func:divergence_concentration) crossed rate_threshold. flagged Bool array, shape (n_chains,). Which chains crossed rate_threshold, independent of whether warn ends up true. num_flagged Number of flagged chains, flagged.sum(). rates Per-chain sampling-phase divergence rate, shape (n_chains,). early_rate, late_rate Per-chain divergence rate in the first and last quarter of the sampling draws, shape (n_chains,). NaN when computed from :func:divergence_concentration_from_counts (no per-draw resolution available). median_other_rate For each chain, the median rates value across the other n_chains - 1 chains, shape (n_chains,). NaN when n_chains <= 1. total_divergences Total divergence count D summed over all chains. num_chains Number of chains M. rate_threshold The threshold that was applied (echoed back for the message / for callers that only keep the report). multinomial_p_value Bonferroni-corrected exchangeable-null tail probability for the single worst chain, min(1, M * P(Binomial(D, 1/M) >= d_max)). Context only; never drives warn. NaN when D=0.

Source code in blackjax/diagnostics.py

class DivergenceConcentrationReport(NamedTuple):
    """Per-run divergence-concentration diagnostic report.

    All fields are plain JAX numerics — no strings — so the statistic
    itself stays JIT-compilable.  Pass the report to
    :func:`format_divergence_warning` to render a human-readable message
    from concrete (non-traced) values.

    Attributes
    ----------
    warn
        Bool. Whether a minority of chains (see the module notes on
        :func:`divergence_concentration`) crossed ``rate_threshold``.
    flagged
        Bool array, shape ``(n_chains,)``. Which chains crossed
        ``rate_threshold``, independent of whether ``warn`` ends up true.
    num_flagged
        Number of flagged chains, ``flagged.sum()``.
    rates
        Per-chain sampling-phase divergence rate, shape ``(n_chains,)``.
    early_rate, late_rate
        Per-chain divergence rate in the first and last quarter of the
        sampling draws, shape ``(n_chains,)``. ``NaN`` when computed from
        :func:`divergence_concentration_from_counts` (no per-draw
        resolution available).
    median_other_rate
        For each chain, the median ``rates`` value across the *other*
        ``n_chains - 1`` chains, shape ``(n_chains,)``. ``NaN`` when
        ``n_chains <= 1``.
    total_divergences
        Total divergence count ``D`` summed over all chains.
    num_chains
        Number of chains ``M``.
    rate_threshold
        The threshold that was applied (echoed back for the message /
        for callers that only keep the report).
    multinomial_p_value
        Bonferroni-corrected exchangeable-null tail probability for the
        single worst chain, ``min(1, M * P(Binomial(D, 1/M) >= d_max))``.
        Context only; never drives ``warn``. ``NaN`` when ``D=0``.
    """

    warn: Array
    flagged: Array
    num_flagged: Array
    rates: Array
    early_rate: Array
    late_rate: Array
    median_other_rate: Array
    total_divergences: Array
    num_chains: Array
    rate_threshold: Array
    multinomial_p_value: Array

Sequential Monte Carlo¶

probjax.inference.smc ¶

smc(*, path=None, **kwargs)

Build an SMC kernel from a path object/callable.

Source code in probjax/inference/smc/__init__.py

def smc(*, path=None, **kwargs):
    """Build an SMC kernel from a path object/callable."""
    from probjax.inference.smc.path import GeometricPath
    from probjax.inference.smc.path_smc import path_smc

    if path is None:
        path = GeometricPath()
    if "logprior_fn" not in kwargs:
        kwargs["logprior_fn"] = None
    if "loglikelihood_fn" not in kwargs:
        kwargs["loglikelihood_fn"] = None
    if hasattr(path, "logdensity_fn"):
        return path_smc(path=path, **kwargs)
    if callable(path):
        return path_smc(path=path, **kwargs)
    raise TypeError("path must define logdensity_fn or be callable.")

probjax.inference.adaptive_smc ¶

Bases: SMCKernelAPI

Tempered SMC that picks each temperature to hit a target ESS.

Moves along the geometric path logprior + t * loglikelihood from t = 0 to 1, choosing every increment so the effective sample size after reweighting is target_ess. Initial MCMC parameters come from kernel.init_params(particles, **mcmc_params).

Parameters:

Name	Description	Default
`logprior_fn`	Log prior.	required
`loglikelihood_fn`	Log likelihood.	required
`path`	Tempering path; must be :class:`GeometricPath`.	required
`mcmc_kernel`	MCMC kernel class used to move the particles, e.g. :class:`hmc` (the class, not a built kernel).	required
`target_ess`	Target effective sample size as a fraction of the particles.	required
`root_solver`	Solver for the next temperature.	required
`num_mcmc_steps`	MCMC steps per SMC iteration.	required
`resampling_fn`	BlackJAX resampling function.	required
`update_strategy`	BlackJAX particle update strategy.	required
`batch_size`	Particles per batch in the MCMC update; `0` updates all at once.	required
`**mcmc_kernel_kwargs`	Options for `mcmc_kernel`, e.g. `num_integration_steps`.	required

Source code in probjax/inference/smc/adaptive.py

class adaptive_smc(SMCKernelAPI):
    """Tempered SMC that picks each temperature to hit a target ESS.

    Moves along the geometric path ``logprior + t * loglikelihood`` from
    ``t = 0`` to ``1``, choosing every increment so the effective sample size
    after reweighting is ``target_ess``. Initial MCMC parameters come from
    ``kernel.init_params(particles, **mcmc_params)``.

    Args:
        logprior_fn: Log prior.
        loglikelihood_fn: Log likelihood.
        path: Tempering path; must be :class:`GeometricPath`.
        mcmc_kernel: MCMC kernel class used to move the particles, e.g.
            :class:`hmc` (the class, not a built kernel).
        target_ess: Target effective sample size as a fraction of the
            particles.
        root_solver: Solver for the next temperature.
        num_mcmc_steps: MCMC steps per SMC iteration.
        resampling_fn: BlackJAX resampling function.
        update_strategy: BlackJAX particle update strategy.
        batch_size: Particles per batch in the MCMC update; ``0`` updates all
            at once.
        **mcmc_kernel_kwargs: Options for ``mcmc_kernel``, e.g.
            ``num_integration_steps``.
    """

    init = staticmethod(init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_step)

probjax.inference.persistent_smc ¶

Bases: SMCKernelAPI

Persistent sampling: tempered SMC that reuses all past particles.

Particles from every earlier temperature stay in the population and are reweighted, which lowers the variance for a fixed budget. Follows a fixed schedule of n_schedule temperatures; initial MCMC parameters come from kernel.init_params(particles, **mcmc_params).

Parameters:

Name	Description	Default
`logprior_fn`	Log prior.	required
`loglikelihood_fn`	Log likelihood.	required
`mcmc_kernel`	MCMC kernel class used to move the particles, e.g. :class:`hmc` (the class, not a built kernel).	required
`num_mcmc_steps`	MCMC steps per SMC iteration.	required
`resampling_fn`	BlackJAX resampling function.	required
`update_strategy`	BlackJAX particle update strategy.	required
`batch_size`	Particles per batch in the MCMC update; `0` updates all at once.	required
`n_schedule`	Number of tempering steps, forwarded to `init`.	required
`**mcmc_kernel_kwargs`	Options for `mcmc_kernel`, e.g. `num_integration_steps`.	required

Source code in probjax/inference/smc/persistent_smc.py

class persistent_smc(SMCKernelAPI):
    """Persistent sampling: tempered SMC that reuses all past particles.

    Particles from every earlier temperature stay in the population and are
    reweighted, which lowers the variance for a fixed budget. Follows a fixed
    schedule of ``n_schedule`` temperatures; initial MCMC parameters come from
    ``kernel.init_params(particles, **mcmc_params)``.

    Args:
        logprior_fn: Log prior.
        loglikelihood_fn: Log likelihood.
        mcmc_kernel: MCMC kernel class used to move the particles, e.g.
            :class:`hmc` (the class, not a built kernel).
        num_mcmc_steps: MCMC steps per SMC iteration.
        resampling_fn: BlackJAX resampling function.
        update_strategy: BlackJAX particle update strategy.
        batch_size: Particles per batch in the MCMC update; ``0`` updates all
            at once.
        n_schedule: Number of tempering steps, forwarded to ``init``.
        **mcmc_kernel_kwargs: Options for ``mcmc_kernel``, e.g.
            ``num_integration_steps``.
    """

    init = staticmethod(init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_step)

probjax.inference.adaptive_persistent_smc ¶

Bases: SMCKernelAPI

Persistent sampling with adaptively chosen temperatures.

Like :class:persistent_smc, but every temperature is chosen so the ESS of the persistent population is target_ess.

Parameters:

Name	Description	Default
`logprior_fn`	Log prior.	required
`loglikelihood_fn`	Log likelihood.	required
`mcmc_kernel`	MCMC kernel class used to move the particles, e.g. :class:`hmc` (the class, not a built kernel).	required
`target_ess`	Target effective sample size as a fraction of the particles.	required
`root_solver`	Solver for the next temperature.	required
`num_mcmc_steps`	MCMC steps per SMC iteration.	required
`resampling_fn`	BlackJAX resampling function.	required
`update_strategy`	BlackJAX particle update strategy.	required
`batch_size`	Particles per batch in the MCMC update; `0` updates all at once.	required
`max_iterations`	Maximum number of tempering steps, forwarded to `init`.	required
`**mcmc_kernel_kwargs`	Options for `mcmc_kernel`, e.g. `num_integration_steps`.	required

Source code in probjax/inference/smc/adaptive_persistent_smc.py

class adaptive_persistent_smc(SMCKernelAPI):
    """Persistent sampling with adaptively chosen temperatures.

    Like :class:`persistent_smc`, but every temperature is chosen so the ESS
    of the persistent population is ``target_ess``.

    Args:
        logprior_fn: Log prior.
        loglikelihood_fn: Log likelihood.
        mcmc_kernel: MCMC kernel class used to move the particles, e.g.
            :class:`hmc` (the class, not a built kernel).
        target_ess: Target effective sample size as a fraction of the
            particles.
        root_solver: Solver for the next temperature.
        num_mcmc_steps: MCMC steps per SMC iteration.
        resampling_fn: BlackJAX resampling function.
        update_strategy: BlackJAX particle update strategy.
        batch_size: Particles per batch in the MCMC update; ``0`` updates all
            at once.
        max_iterations: Maximum number of tempering steps, forwarded to
            ``init``.
        **mcmc_kernel_kwargs: Options for ``mcmc_kernel``, e.g.
            ``num_integration_steps``.
    """

    init = staticmethod(init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_step)

probjax.inference.path_smc ¶

Bases: SMCKernelAPI

SMC along a user-defined path of intermediate targets.

The runner supplies the path parameter for each step (e.g. a temperature schedule); path turns it into a log density and incremental weights. Usually built through :func:smc. Initial MCMC parameters come from kernel.init_params(particles, **mcmc_params).

Parameters:

Name	Description	Default
`logprior_fn`	Log prior.	required
`loglikelihood_fn`	Log likelihood.	required
`path`	Path object, e.g. :class:`GeometricPath` or :class:`PartialPosteriorsPath`.	required
`mcmc_kernel`	MCMC kernel class used to move the particles, e.g. :class:`hmc` (the class, not a built kernel).	required
`num_mcmc_steps`	MCMC steps per SMC iteration.	required
`resampling_fn`	BlackJAX resampling function.	required
`update_strategy`	BlackJAX particle update strategy.	required
`path_kwargs`	Extra arguments for `path`; inferred from `mcmc_kernel_kwargs` when `None`.	required
`batch_size`	Particles per batch in the MCMC update; `0` updates all at once.	required
`**mcmc_kernel_kwargs`	Options for `mcmc_kernel`, e.g. `num_integration_steps`.	required

Source code in probjax/inference/smc/path_smc.py

class path_smc(SMCKernelAPI):
    """SMC along a user-defined path of intermediate targets.

    The runner supplies the path parameter for each step (e.g. a temperature
    schedule); ``path`` turns it into a log density and incremental weights.
    Usually built through :func:`smc`. Initial MCMC parameters come from
    ``kernel.init_params(particles, **mcmc_params)``.

    Args:
        logprior_fn: Log prior.
        loglikelihood_fn: Log likelihood.
        path: Path object, e.g. :class:`GeometricPath` or
            :class:`PartialPosteriorsPath`.
        mcmc_kernel: MCMC kernel class used to move the particles, e.g.
            :class:`hmc` (the class, not a built kernel).
        num_mcmc_steps: MCMC steps per SMC iteration.
        resampling_fn: BlackJAX resampling function.
        update_strategy: BlackJAX particle update strategy.
        path_kwargs: Extra arguments for ``path``; inferred from
            ``mcmc_kernel_kwargs`` when ``None``.
        batch_size: Particles per batch in the MCMC update; ``0`` updates all
            at once.
        **mcmc_kernel_kwargs: Options for ``mcmc_kernel``, e.g.
            ``num_integration_steps``.
    """

    init = staticmethod(init)
    init_params = staticmethod(init_params)
    build_step = staticmethod(build_step)

probjax.inference.GeometricPath `dataclass` ¶

Geometric (tempered) path: logprior + t * loglikelihood.

Source code in probjax/inference/smc/path.py

@dataclass(frozen=True)
class GeometricPath:
    """Geometric (tempered) path: logprior + t * loglikelihood."""

    is_geometric: bool = True

    def initial_param(self, **_):
        return 0.0

    def logdensity_fn(
        self, tempering_param, *, logprior_fn: Callable, loglikelihood_fn: Callable, **_
    ):
        def logdensity(x):
            return logprior_fn(x) + tempering_param * loglikelihood_fn(x)

        return logdensity

probjax.inference.PartialPosteriorsPath `dataclass` ¶

Partial posterior (data tempering) path.

Source code in probjax/inference/smc/path.py

@dataclass(frozen=True)
class PartialPosteriorsPath:
    """Partial posterior (data tempering) path."""

    is_geometric: bool = False

    def initial_param(self, *, num_datapoints: int, **_):
        return jnp.zeros(num_datapoints)

    def logdensity_fn(
        self,
        tempering_param,
        *,
        partial_logposterior_factory: Callable,
        **_,
    ):
        return partial_logposterior_factory(tempering_param)

SMC extensions¶

See the SMC guide for BlackJAX presets and adaptive execution.

probjax.inference.waste_free_strategy ¶

waste_free_strategy(num_particles, p)

BlackJAX waste-free update: retain p states per chain.

num_particles must be divisible by p. When supplying this strategy manually, set num_mcmc_steps=None; p determines the chain length instead.

Source code in probjax/inference/smc/ports.py

def waste_free_strategy(num_particles, p):
    """BlackJAX waste-free update: retain p states per chain.

    num_particles must be divisible by p. When supplying this strategy manually,
    set num_mcmc_steps=None; p determines the chain length instead.
    """
    if p < 1 or num_particles < 1:
        raise ValueError("num_particles and p must be positive.")
    return _waste_free_strategy(num_particles, p)

probjax.inference.waste_free_smc ¶

waste_free_smc(logprior_fn, loglikelihood_fn, *, num_particles, p=4, adaptive=False, **kwargs)

Convenient waste-free preset for fixed or adaptive geometric/path SMC.

Source code in probjax/inference/smc/ports.py

def waste_free_smc(
    logprior_fn, loglikelihood_fn, *, num_particles, p=4, adaptive=False, **kwargs
):
    """Convenient waste-free preset for fixed or adaptive geometric/path SMC."""
    # Deferred: this module is re-exported by smc/__init__, so importing the
    # package-level factories at top level would be a self-parent cycle.
    from probjax.inference.smc import adaptive_smc_kernel, smc

    if 'num_mcmc_steps' in kwargs or 'update_strategy' in kwargs:
        raise ValueError("Use p to configure waste-free chain length.")
    constructor = adaptive_smc_kernel if adaptive else smc
    return constructor(
        logprior_fn=logprior_fn,
        loglikelihood_fn=loglikelihood_fn,
        num_mcmc_steps=None,
        update_strategy=waste_free_strategy(num_particles, p),
        **kwargs,
    )

probjax.inference.tuned_smc ¶

tuned_smc(logprior_fn, loglikelihood_fn, *, mcmc_kernel, mcmc_parameters, parameter_update_fn, adaptive=True, num_mcmc_steps=10, target_ess=0.8, batch_size=0, resampling_fn=systematic, **mcmc_kernel_kwargs)

Wrap BlackJAX inner-kernel tuning with a probjax MCMC kernel.

parameter_update_fn(key, new_smc_state, info) returns the next MCMC parameter dictionary, with BlackJAX's leading shared (1) or per-particle (N) axis. mcmc_parameters supplies initial shared, unbatched parameter values. adaptive=False accepts explicit temperatures through SMC.run; otherwise use SMC.run_adaptive. Parameters live in state.parameter_override.

Source code in probjax/inference/smc/ports.py

def tuned_smc(
    logprior_fn,
    loglikelihood_fn,
    *,
    mcmc_kernel,
    mcmc_parameters,
    parameter_update_fn,
    adaptive=True,
    num_mcmc_steps=10,
    target_ess=0.8,
    batch_size=0,
    resampling_fn=systematic,
    **mcmc_kernel_kwargs,
):
    """Wrap BlackJAX inner-kernel tuning with a probjax MCMC kernel.

    parameter_update_fn(key, new_smc_state, info) returns the next MCMC parameter
    dictionary, with BlackJAX's leading shared (1) or per-particle (N) axis.
    mcmc_parameters supplies initial *shared*, unbatched parameter values.
    adaptive=False accepts explicit temperatures through SMC.run; otherwise use
    SMC.run_adaptive. Parameters live in state.parameter_override.
    """
    init_mcmc, step_mcmc = make_mcmc_adapter(mcmc_kernel, **mcmc_kernel_kwargs)
    parameters = extend_params(mcmc_parameters)
    algorithm = blackjax.adaptive_tempered_smc if adaptive else blackjax.tempered_smc
    kwargs = dict(batch_size=batch_size)
    if adaptive:
        kwargs['target_ess'] = target_ess
    delegate = inner_kernel_tuning.as_top_level_api(
        algorithm,
        logprior_fn,
        loglikelihood_fn,
        step_mcmc,
        init_mcmc,
        resampling_fn,
        parameter_update_fn,
        parameters,
        num_mcmc_steps=num_mcmc_steps,
        **kwargs,
    )

    def step(key, state, mcmc_parameters=None, tempering_param=None):
        if adaptive:
            return delegate.step(key, state)
        return delegate.step(key, state, tempering_param=tempering_param)

    return Kernel(delegate.init, step, lambda *a, **kw: parameters)

probjax.inference.pretuned_smc ¶

pretuned_smc(logprior_fn, loglikelihood_fn, *, mcmc_kernel, mcmc_parameters, num_particles, sigma_parameters, alpha=0.5, adaptive=True, num_mcmc_steps=10, target_ess=0.8, batch_size=0, positive_parameters=None, natural_parameters=None, performance_of_chain_measure_factory=default_measure_factory, resampling_fn=systematic, **mcmc_kernel_kwargs)

BlackJAX pilot-move pretuning, retaining a distribution of move parameters.

sigma_parameters selects parameters to perturb and their noise scales. Initial mcmc_parameters use BlackJAX's explicit leading 1/N batch convention (unlike tuned_smc). The default ESJD measure requires a full shared inverse mass matrix of shape (1,D,D). Supply a custom measure factory for other moves. Pilot and production transitions use independent keys. Production moves respect batch_size; BlackJAX's pilot currently evaluates its population together.

Source code in probjax/inference/smc/ports.py

def pretuned_smc(
    logprior_fn,
    loglikelihood_fn,
    *,
    mcmc_kernel,
    mcmc_parameters,
    num_particles,
    sigma_parameters,
    alpha=0.5,
    adaptive=True,
    num_mcmc_steps=10,
    target_ess=0.8,
    batch_size=0,
    positive_parameters=None,
    natural_parameters=None,
    performance_of_chain_measure_factory=pretuning.default_measure_factory,
    resampling_fn=systematic,
    **mcmc_kernel_kwargs,
):
    """BlackJAX pilot-move pretuning, retaining a distribution of move parameters.

    sigma_parameters selects parameters to perturb and their noise scales.
    Initial mcmc_parameters use BlackJAX's explicit leading 1/N batch convention
    (unlike tuned_smc). The default ESJD measure requires a full shared inverse
    mass matrix of shape (1,D,D). Supply a custom measure factory for other moves.
    Pilot and production transitions use independent keys. Production moves
    respect batch_size; BlackJAX's pilot currently evaluates its population together.
    """
    init_mcmc, step_mcmc = make_mcmc_adapter(mcmc_kernel, **mcmc_kernel_kwargs)
    pilot = pretuning.build_pretune(
        init_mcmc,
        step_mcmc,
        alpha,
        sigma_parameters,
        num_particles,
        performance_of_chain_measure_factory=performance_of_chain_measure_factory,
        positive_parameters=positive_parameters,
        natural_parameters=natural_parameters,
    )
    move = from_mcmc.build_kernel(
        step_mcmc, init_mcmc, resampling_fn, batch_size=batch_size
    )

    def update(key, state, num_steps, parameters, logdensity, logweights):
        pilot_key, move_key = jax.random.split(key)
        parameters = pilot(
            pilot_key,
            inner_kernel_tuning.StateWithParameterOverride(state, dict(parameters)),
            logdensity,
        )
        state, info = move(
            move_key, state, num_steps, parameters, logdensity, logweights
        )
        return state, pretuning.SMCInfoWithParameterDistribution(info, parameters)

    algorithm = blackjax.adaptive_tempered_smc if adaptive else blackjax.tempered_smc
    kwargs = dict(update_particles_fn=update, batch_size=batch_size)
    if adaptive:
        kwargs['target_ess'] = target_ess

    def init(particles):
        if jax.tree.leaves(particles)[0].shape[0] != num_particles:
            raise ValueError("num_particles must match the initial population.")
        return inner_kernel_tuning.StateWithParameterOverride(
            blackjax.tempered_smc.init(particles), dict(mcmc_parameters)
        )

    def step(key, state, mcmc_parameters=None, tempering_param=None):
        delegate = algorithm(
            logprior_fn,
            loglikelihood_fn,
            step_mcmc,
            init_mcmc,
            state.parameter_override,
            resampling_fn,
            num_mcmc_steps=num_mcmc_steps,
            **kwargs,
        )
        step_kwargs = {} if adaptive else dict(tempering_param=tempering_param)
        sampler, info = delegate.step(key, state.sampler_state, **step_kwargs)
        return inner_kernel_tuning.StateWithParameterOverride(
            sampler, info.parameter_override
        ), info.smc_info

    return Kernel(init, step, lambda *a, **kw: dict(mcmc_parameters))

probjax.inference.population_diagnostics ¶

population_diagnostics(particles, log_weights, ancestors=None)

Weight ESS and exact population/ancestral diversity (O(N log N) sorting).

Diagnostics do not imply independent posterior samples. ancestors may be composed origin labels for measuring collapse over many resampling steps.

Source code in probjax/inference/smc/temporal_utils.py

def population_diagnostics(particles, log_weights, ancestors=None):
    """Weight ESS and exact population/ancestral diversity (O(N log N) sorting).

    Diagnostics do not imply independent posterior samples. ancestors may be
    composed origin labels for measuring collapse over many resampling steps.
    """
    x = population_matrix(particles)
    order = jnp.lexsort(x.T[::-1])
    unique = 1 + jnp.sum(jnp.any(jnp.diff(x[order], axis=0) != 0, axis=1))
    if ancestors is None:
        unique_ancestors = jnp.array(-1)  # unavailable
    else:
        labels = jnp.sort(ancestors)
        unique_ancestors = 1 + jnp.sum(jnp.diff(labels) != 0)
    return PopulationDiagnostics(
        ess(log_weights), unique, unique_ancestors, jax.nn.softmax(log_weights).max()
    )

probjax.inference.summarize_replicates ¶

summarize_replicates(estimates)

Mean and Monte Carlo standard error across independent RUN estimates.

Supply posterior expectation estimates (or evidence estimates on the desired scale), not individual correlated particles from one run.

Source code in probjax/inference/smc/temporal_utils.py

def summarize_replicates(estimates):
    """Mean and Monte Carlo standard error across independent RUN estimates.

    Supply posterior expectation estimates (or evidence estimates on the desired
    scale), not individual correlated particles from one run.
    """
    values = jnp.asarray(estimates)
    if values.ndim < 1 or values.shape[0] < 2:
        raise ValueError("At least two independent run estimates are required.")
    return ReplicateDiagnostics(
        values.mean(0),
        values.std(0, ddof=1) / jnp.sqrt(values.shape[0]),
        jnp.array(values.shape[0]),
    )

Twisted SMC for diffusion models¶

probjax.inference.run_pretrained_diffusion_smc ¶

run_pretrained_diffusion_smc(key, *, model, event_spec=None, log_potential_fn, num_particles=256, num_steps=None, t_min=None, t_max=None, temperature=1.0, ess_threshold=0.5, twist=None, inputs=None)

Run twisted SMC steering a frozen diffusion model toward a potential.

Parameters:

Name	Type	Description	Default
`key`	`RngKey`	RNG key.	required
`model`		diffusion model exposing `solver_cfg.solve_schedule`, `_build_sde_drift_and_diffusion(inputs)`, `denoise(t, x, **inputs)` and `_sample_base(rng, sample_shape, spec, t_max=...)` (i.e. a :class:`DiffusionDenoiser` or a duck-typed equivalent).	required
`event_spec`		single-event shape/spec for sampling. Defaults to the model's `event_spec`.	`None`
`log_potential_fn`	`Callable[[PyTree[Array]], Array]`	single-particle `x_0 -> scalar` potential `V`. The target is `p_phi(x_0) * exp(temperature * V(x_0))`.	required
`num_particles`	`int`	particle count `N`.	`256`
`num_steps`	`Optional[int]`	reverse steps; defaults to `solver_cfg.num_steps`.	`None`
`t_min/t_max`		time range; default to `train_cfg` bounds.	required
`temperature`	`float`	scales `V` (`0.0` recovers the prior).	`1.0`
`ess_threshold`	`float`	resample when `ESS / N` falls below this.	`0.5`
`twist`	`Optional[TwistFn]`	custom :class:`TwistFn`. Defaults to :class:`DenoisedPotentialTwist` around `model.denoise`.	`None`
`inputs`	`Optional[dict]`	named model inputs (`context`, ...) forwarded to drift/diffusion/denoise.	`None`

Returns:

Type	Description
`SMCResult`	class:`SMCResult` with `state` a :class:`PretrainedTwistState`
`SMCResult`	of clean particles and normalized weights, and `log_evidence`
`SMCResult`	estimating `log E_prior[exp(temperature * V)]`.

Source code in probjax/inference/smc/pretrained_twist.py

def run_pretrained_diffusion_smc(
    key: RngKey,
    *,
    model,
    event_spec=None,
    log_potential_fn: Callable[[PyTree[Array]], Array],
    num_particles: int = 256,
    num_steps: Optional[int] = None,
    t_min: Optional[float] = None,
    t_max: Optional[float] = None,
    temperature: float = 1.0,
    ess_threshold: float = 0.5,
    twist: Optional[TwistFn] = None,
    inputs: Optional[dict] = None,
) -> SMCResult:
    """Run twisted SMC steering a frozen diffusion model toward a potential.

    Args:
        key: RNG key.
        model: diffusion model exposing ``solver_cfg.solve_schedule``,
            ``_build_sde_drift_and_diffusion(inputs)``,
            ``denoise(t, x, **inputs)`` and
            ``_sample_base(rng, sample_shape, spec, t_max=...)`` (i.e. a
            :class:`DiffusionDenoiser` or a duck-typed equivalent).
        event_spec: single-event shape/spec for sampling. Defaults to the
            model's ``event_spec``.
        log_potential_fn: single-particle ``x_0 -> scalar`` potential ``V``.
            The target is ``p_phi(x_0) * exp(temperature * V(x_0))``.
        num_particles: particle count ``N``.
        num_steps: reverse steps; defaults to ``solver_cfg.num_steps``.
        t_min/t_max: time range; default to ``train_cfg`` bounds.
        temperature: scales ``V`` (``0.0`` recovers the prior).
        ess_threshold: resample when ``ESS / N`` falls below this.
        twist: custom :class:`TwistFn`. Defaults to
            :class:`DenoisedPotentialTwist` around ``model.denoise``.
        inputs: named model inputs (``context``, ...) forwarded to
            drift/diffusion/denoise.

    Returns:
        :class:`SMCResult` with ``state`` a :class:`PretrainedTwistState`
        of clean particles and normalized weights, and ``log_evidence``
        estimating ``log E_prior[exp(temperature * V)]``.
    """
    if num_particles < 1:
        raise ValueError("num_particles must be positive.")
    if not 0.0 < ess_threshold <= 1.0:
        raise ValueError("ess_threshold must be in (0, 1].")
    inputs = dict(inputs or {})

    cfg = model.solver_cfg
    if cfg is None:
        raise ValueError("model.solver_cfg is not set.")
    train_cfg = model.train_cfg
    if t_min is None:
        t_min = float(train_cfg.t_min)
    if t_max is None:
        t_max = float(train_cfg.t_max)
    steps = cfg.num_steps if num_steps is None else int(num_steps)
    if steps < 1:
        raise ValueError("num_steps must be positive.")

    ts = jnp.asarray(cfg.solve_schedule(t_min=t_min, t_max=t_max, num_steps=steps))
    if ts.ndim != 1 or ts.shape[0] < 2:
        raise ValueError("solve_schedule must return at least 2 time points.")
    if ts[1] > ts[0]:  # ascend -> flip so we integrate t_max -> t_min
        ts = ts[::-1]

    drift_fn, diffusion_fn = model._build_sde_drift_and_diffusion(inputs)

    def denoise_fn(t: Array, x: PyTree[Array]) -> PyTree[Array]:
        return model.denoise(t, x, **inputs)

    if twist is None:
        twist = DenoisedPotentialTwist(
            denoise_fn=denoise_fn,
            log_potential_fn=log_potential_fn,
            temperature=temperature,
        )
    potential_batch = jax.vmap(log_potential_fn)

    if event_spec is None:
        event_spec = model.event_spec
        if event_spec is None:
            raise TypeError(
                "event_spec is required: pass it explicitly or set a default "
                "on the model via set_event_spec()."
            )
    spec = normalize_spec(event_spec)

    key, key_init = jax.random.split(key)
    particles = model._sample_base(key_init, (num_particles,), spec, t_max=t_max)

    log_psi = jnp.asarray(twist.log_twist(ts[0], particles), dtype=jnp.float32)
    # Initial weight w_T = gamma_T / q = psi_T (proposal equals prior noise).
    log_z = jax.scipy.special.logsumexp(log_psi) - jnp.log(num_particles)
    log_weights = normalize_log_weights(log_psi)

    ess_trace, resampled_trace, increment_trace, ancestor_trace = [], [], [], []

    n_transitions = ts.shape[0] - 1
    for i in range(n_transitions):
        t_cur, t_next = ts[i], ts[i + 1]
        h = float(t_cur - t_next)
        if h <= 0:
            raise ValueError("time grid must strictly decrease.")
        key, key_noise, key_resample = jax.random.split(key, 3)

        drift = drift_fn(t_cur, particles)
        diffusion = diffusion_fn(t_cur, particles)
        noise = random_like(key_noise, particles)
        scale = jnp.sqrt(h)

        def _em_step(x, d, g, e, h=h, scale=scale):
            return x - d * h + g * scale * e

        particles_new = jax.tree.map(
            _em_step,
            particles,
            drift,
            diffusion,
            noise,
        )

        if i == n_transitions - 1:  # exact terminal potential on clean x_0
            log_psi_new = temperature * jnp.asarray(
                potential_batch(particles_new), dtype=jnp.float32
            )
        else:
            log_psi_new = jnp.asarray(
                twist.log_twist(t_next, particles_new), dtype=jnp.float32
            )

        log_unnorm = log_weights + (log_psi_new - log_psi)
        increment = log_weighted_mean_exp(log_weights, log_psi_new - log_psi)
        log_z = log_z + increment
        log_weights = normalize_log_weights(log_unnorm)

        current_ess = ess(log_weights)
        do_resample = bool(current_ess < ess_threshold * num_particles)
        if num_particles == 1:
            do_resample = False
        if do_resample:
            idx = systematic(key_resample, jnp.exp(log_weights), num_particles).astype(
                jnp.int32
            )
            particles = take_batch(particles_new, idx)
            log_weights = jnp.full((num_particles,), -jnp.log(num_particles))
            log_psi = log_psi_new[idx]
        else:
            particles = particles_new
            idx = jnp.arange(num_particles, dtype=jnp.int32)
            log_psi = log_psi_new

        ess_trace.append(current_ess)
        resampled_trace.append(do_resample)
        increment_trace.append(increment)
        ancestor_trace.append(idx)

    state = PretrainedTwistState(particles=particles, log_weights=log_weights)
    info = PretrainedTwistInfo(
        ess=jnp.asarray(ess_trace),
        resampled=jnp.asarray(resampled_trace),
        log_z_increments=jnp.asarray(increment_trace),
        ancestors=jnp.stack(ancestor_trace),
    )
    return SMCResult(
        state=state,
        params=None,
        info=info,
        log_evidence=log_z,
        final_info=None,
        num_steps=n_transitions,
        completed=True,
        key=key,
    )

probjax.inference.PretrainedTwistState ¶

Bases: NamedTuple

Source code in probjax/inference/smc/pretrained_twist.py

class PretrainedTwistState(NamedTuple):
    particles: PyTree[Array]
    log_weights: Array  # normalized, shape (num_particles,)

probjax.inference.PretrainedTwistInfo ¶

Bases: NamedTuple

Source code in probjax/inference/smc/pretrained_twist.py

class PretrainedTwistInfo(NamedTuple):
    ess: Array  # (num_transitions,) ESS before resampling each step
    resampled: Array  # (num_transitions,) bool
    log_z_increments: Array  # (num_transitions,) evidence increments
    ancestors: Any  # (num_transitions, num_particles) resample indices

probjax.inference.run_rdsmc ¶

run_rdsmc(key, *, log_target_fn, event_spec, num_particles=256, num_steps=32, num_inner=64, schedule=None, t_min=0.001, t_max=1.0, max_score_norm=50.0, ess_threshold=0.5)

Sample an unnormalized target with reverse-diffusion SMC (no training).

Parameters:

Name	Type	Description	Default
`key`	`RngKey`	RNG key.	required
`log_target_fn`	`Callable[[PyTree[Array]], Array]`	single-particle unnormalized `log pi(x_0)`.	required
`event_spec`		single-event shape/spec; base noise is standard normal.	required
`num_particles`	`int`	outer particle count `N`.	`256`
`num_steps`	`int`	reverse time steps `T`.	`32`
`num_inner`	`int`	inner importance samples `M` per score estimate.	`64`
`schedule`		forward :class:`NoiseScheduleProtocol` (default VP). Only `scale/std/drift/diffusion` are used.	`None`
`t_min/t_max`		reverse time range. `t_min` must stay slightly above zero so the denoising posterior is non-degenerate.	required
`max_score_norm`	`float \| None`	clip for MC score norms (paper B.3.2).	`50.0`
`ess_threshold`	`float`	resample when `ESS / N` falls below this.	`0.5`

Returns:

Type	Description
`SMCResult`	class:`SMCResult` with `state` an :class:`RDSMCState` of
`SMCResult`	approximate target samples and `log_evidence` an unbiased
`SMCResult`	estimate of `log Z`.

Source code in probjax/inference/smc/rdsmc.py

def run_rdsmc(
    key: RngKey,
    *,
    log_target_fn: Callable[[PyTree[Array]], Array],
    event_spec,
    num_particles: int = 256,
    num_steps: int = 32,
    num_inner: int = 64,
    schedule=None,
    t_min: float = 1e-3,
    t_max: float = 1.0,
    max_score_norm: float | None = 50.0,
    ess_threshold: float = 0.5,
) -> SMCResult:
    """Sample an unnormalized target with reverse-diffusion SMC (no training).

    Args:
        key: RNG key.
        log_target_fn: single-particle unnormalized ``log pi(x_0)``.
        event_spec: single-event shape/spec; base noise is standard normal.
        num_particles: outer particle count ``N``.
        num_steps: reverse time steps ``T``.
        num_inner: inner importance samples ``M`` per score estimate.
        schedule: forward :class:`NoiseScheduleProtocol` (default VP).
            Only ``scale/std/drift/diffusion`` are used.
        t_min/t_max: reverse time range. ``t_min`` must stay slightly
            above zero so the denoising posterior is non-degenerate.
        max_score_norm: clip for MC score norms (paper B.3.2).
        ess_threshold: resample when ``ESS / N`` falls below this.

    Returns:
        :class:`SMCResult` with ``state`` an :class:`RDSMCState` of
        approximate target samples and ``log_evidence`` an unbiased
        estimate of ``log Z``.
    """
    if num_particles < 1:
        raise ValueError("num_particles must be positive.")
    if num_steps < 1:
        raise ValueError("num_steps must be positive.")
    if num_inner < 1:
        raise ValueError("num_inner must be positive.")
    if not 0.0 < ess_threshold <= 1.0:
        raise ValueError("ess_threshold must be in (0, 1].")
    if not 0.0 < t_min < t_max:
        raise ValueError("require 0 < t_min < t_max.")

    if schedule is None:
        schedule = VPNoiseSchedule(t_min=t_min, t_max=t_max)
    solver = BaseSolverConfig(schedule=schedule, num_steps=num_steps)
    ts = jnp.asarray(solver.solve_schedule(t_min=t_min, t_max=t_max))
    if ts.ndim != 1 or ts.shape[0] < 2:
        raise ValueError("solve_schedule must return at least 2 time points.")
    if ts[1] > ts[0]:
        ts = ts[::-1]

    spec = normalize_spec(event_spec)
    potential_batch = jax.vmap(log_target_fn)

    def estimate_batch(est_keys: Array, t: Array, batch: PyTree[Array]):
        def _single(k: Array, x: PyTree[Array]):
            return is_score_and_marginal(
                k,
                x,
                t,
                log_target_fn=log_target_fn,
                alpha=schedule.scale(t),
                sigma=schedule.std(t),
                num_inner=num_inner,
                max_score_norm=max_score_norm,
            )

        return jax.vmap(_single)(est_keys, batch)

    key, key_init, key_est = jax.random.split(key, 3)
    particles = sample_normal(key_init, (num_particles,), spec)
    est_keys = jax.random.split(key_est, num_particles)
    scores, log_psi, est_info = estimate_batch(est_keys, ts[0], particles)
    inner_ess_trace = [jnp.mean(est_info.ess)]

    # Initial weight pi_hat_T / q with standard-normal base q.
    log_unnorm = log_psi - _standard_normal_logpdf(particles)
    log_z = jax.scipy.special.logsumexp(log_unnorm) - jnp.log(num_particles)
    log_weights = normalize_log_weights(log_unnorm)

    ess_trace, resampled_trace, increment_trace, ancestor_trace = [], [], [], []
    n_transitions = ts.shape[0] - 1
    for i in range(n_transitions):
        t_cur, t_next = ts[i], ts[i + 1]
        h = float(t_cur - t_next)
        if h <= 0:
            raise ValueError("time grid must strictly decrease.")
        key, key_noise, key_est, key_resample = jax.random.split(key, 4)

        fwd_drift = schedule.drift(t_cur, particles)
        diffusion = schedule.diffusion(t_cur, particles)
        rev_drift = jax.tree.map(
            lambda f, g, s: f - g**2 * s, fwd_drift, diffusion, scores
        )
        noise = random_like(key_noise, particles)
        scale = jnp.sqrt(h)

        def _em_step(x, d, g, e, h=h, scale=scale):
            return x - d * h + g * scale * e

        particles_new = jax.tree.map(_em_step, particles, rev_drift, diffusion, noise)

        est_keys = jax.random.split(key_est, num_particles)
        scores_new, log_psi_new, est_info = estimate_batch(
            est_keys, t_next, particles_new
        )
        inner_ess_trace.append(jnp.mean(est_info.ess))

        log_unnorm = log_weights + (log_psi_new - log_psi)
        increment = log_weighted_mean_exp(log_weights, log_psi_new - log_psi)
        log_z = log_z + increment
        log_weights = normalize_log_weights(log_unnorm)

        current_ess = ess(log_weights)
        do_resample = bool(current_ess < ess_threshold * num_particles)
        if num_particles == 1:
            do_resample = False
        if do_resample:
            idx = systematic(key_resample, jnp.exp(log_weights), num_particles).astype(
                jnp.int32
            )
            particles = take_batch(particles_new, idx)
            scores = take_batch(scores_new, idx)
            log_psi = log_psi_new[idx]
            log_weights = jnp.full((num_particles,), -jnp.log(num_particles))
        else:
            particles = particles_new
            scores = scores_new
            log_psi = log_psi_new
            idx = jnp.arange(num_particles, dtype=jnp.int32)

        ess_trace.append(current_ess)
        resampled_trace.append(do_resample)
        increment_trace.append(increment)
        ancestor_trace.append(idx)

    # Exact terminal correction: replace the t_min estimate with pi_tilde(x_0).
    log_exact = jnp.asarray(potential_batch(particles), dtype=jnp.float32)
    final_increment = log_weighted_mean_exp(log_weights, log_exact - log_psi)
    log_z = log_z + final_increment
    log_weights = normalize_log_weights(log_weights + (log_exact - log_psi))

    state = RDSMCState(particles=particles, log_weights=log_weights)
    info = RDSMCInfo(
        ess=jnp.asarray(ess_trace),
        resampled=jnp.asarray(resampled_trace),
        log_z_increments=jnp.asarray(increment_trace),
        inner_ess=jnp.asarray(inner_ess_trace),
        ancestors=jnp.stack(ancestor_trace),
        final_log_z_increment=final_increment,
    )
    return SMCResult(
        state=state,
        params=None,
        info=info,
        log_evidence=log_z,
        final_info=None,
        num_steps=n_transitions,
        completed=True,
        key=key,
    )

probjax.inference.RDSMCState ¶

Bases: NamedTuple

Source code in probjax/inference/smc/rdsmc.py

class RDSMCState(NamedTuple):
    particles: PyTree[Array]
    log_weights: Array  # normalized, shape (num_particles,)

probjax.inference.RDSMCInfo ¶

Bases: NamedTuple

Source code in probjax/inference/smc/rdsmc.py

class RDSMCInfo(NamedTuple):
    ess: Array  # (num_transitions,) ESS before resampling each step
    resampled: Array  # (num_transitions,) bool
    log_z_increments: Array  # (num_transitions,) evidence increments
    inner_ess: Array  # (num_transitions + 1,) mean inner ESS per estimator call
    ancestors: Any  # (num_transitions, num_particles) resample indices
    final_log_z_increment: Array  # exact-correction evidence increment

Temporal SMC and trajectory inference¶

Temporal inference advances through physical observation times. See the temporal SMC guide and the executable example notebook.

probjax.inference.StateSpaceModel ¶

Bases: NamedTuple

A state-space model written for a single particle.

All functions act on one latent state x, which may be any pytree; the filter vectorizes them over particles.

Attributes:

Name	Type	Description
`init`	`Callable`	`(key, theta) -> x`, a draw from the initial law at `t0`.
`step`	`Callable`	`(key, theta, x, t_previous, t_next) -> x`, a draw from the transition over one interval.
`log_likelihood`	`Callable`	`(theta, x, y, t) -> scalar`, the normalized observation log density.
`transition_logdensity`	`Optional[Callable]`	`(theta, x_next, x_previous, t_previous, t_next) -> scalar`. Optional; needed only for backward path sampling and particle Gibbs.

Source code in probjax/inference/filtering/state_space.py

class StateSpaceModel(NamedTuple):
    """A state-space model written for a single particle.

    All functions act on one latent state ``x``, which may be any pytree; the
    filter vectorizes them over particles.

    Attributes:
        init: ``(key, theta) -> x``, a draw from the initial law at ``t0``.
        step: ``(key, theta, x, t_previous, t_next) -> x``, a draw from the
            transition over one interval.
        log_likelihood: ``(theta, x, y, t) -> scalar``, the normalized
            observation log density.
        transition_logdensity: ``(theta, x_next, x_previous, t_previous, t_next)
            -> scalar``. Optional; needed only for backward path sampling and
            particle Gibbs.
    """

    init: Callable
    step: Callable
    log_likelihood: Callable
    transition_logdensity: Optional[Callable] = None

probjax.inference.LinearGaussianModel ¶

Bases: NamedTuple

A linear-Gaussian state-space model, filtered exactly by a Kalman filter.

Attributes:

Name	Type	Description
`initial`	`Callable`	`theta -> (mean, covariance)` of the state at `t0`.
`transition`	`Callable`	`(theta, t_previous, t_next) -> (Phi, Q)`.
`observation`	`Callable`	`(theta, t) -> (C, R)`.

Source code in probjax/inference/filtering/state_space.py

class LinearGaussianModel(NamedTuple):
    """A linear-Gaussian state-space model, filtered exactly by a Kalman filter.

    Attributes:
        initial: ``theta -> (mean, covariance)`` of the state at ``t0``.
        transition: ``(theta, t_previous, t_next) -> (Phi, Q)``.
        observation: ``(theta, t) -> (C, R)``.
    """

    initial: Callable
    transition: Callable
    observation: Callable

probjax.inference.run_temporal_filter ¶

run_temporal_filter(key, model, theta, ts, observations, *, t0=0.0, num_particles=None, observed=None, history='full')

Filter model at fixed theta over observations at ts after t0.

model is a :class:StateSpaceModel (needs num_particles), a :class:LinearGaussianModel or a :class:TemporalFilter. history is 'full', 'none' or a window size: a positive integer keeps a ring buffer of the last intervals. observed masks missing observations, whose values may be NaN.

Source code in probjax/inference/filtering/temporal.py

def run_temporal_filter(
    key,
    model,
    theta,
    ts,
    observations,
    *,
    t0=0.0,
    num_particles=None,
    observed=None,
    history="full",
):
    """Filter ``model`` at fixed ``theta`` over observations at ``ts`` after ``t0``.

    ``model`` is a :class:`StateSpaceModel` (needs ``num_particles``), a
    :class:`LinearGaussianModel` or a :class:`TemporalFilter`. ``history`` is
    ``'full'``, ``'none'`` or a window size: a positive integer keeps a ring
    buffer of the last intervals. ``observed`` masks missing observations, whose
    values may be NaN.
    """
    ts, observations, observed = validate_observations(ts, observations, observed)
    filt = temporal_filter(model, num_particles)
    key, init_key = jax.random.split(key)
    state = filt.init(init_key, theta, t0)

    def step(key, state, data):
        t, y, mask = data
        return filt.step(key, state, theta, t, y, mask)

    return scan_history(step, key, state, ts, (ts, observations, observed), history)

probjax.inference.temporal_filter ¶

temporal_filter(model, num_particles=None, ess_threshold=0.5)

Build the filter for model: Kalman for linear-Gaussian, else particles.

A :class:TemporalFilter is returned unchanged.

Source code in probjax/inference/filtering/temporal.py

def temporal_filter(model, num_particles=None, ess_threshold=0.5):
    """Build the filter for ``model``: Kalman for linear-Gaussian, else particles.

    A :class:`TemporalFilter` is returned unchanged.
    """
    if isinstance(model, TemporalFilter):
        return model
    if isinstance(model, LinearGaussianModel):
        return _kalman_filter(model)
    if num_particles is None:
        raise ValueError("A StateSpaceModel needs num_particles.")
    return _particle_filter(model, num_particles, ess_threshold)

probjax.inference.temporal_smc ¶

temporal_smc(model, logprior_fn, *, num_particles=None, ess_threshold=0.5, num_moves=1, proposal_fn=None, batch_size=0)

SMC over parameters that assimilates one observation per step.

Each parameter particle carries its own filter: a particle filter for a :class:StateSpaceModel (SMC², with num_particles latent particles) or a Kalman filter for a :class:LinearGaussianModel (exact). After every observation the outer weights are multiplied by the filters' likelihood increments. When the ESS drops below ess_threshold, parameters are resampled together with their filters and moved by num_moves pseudo-marginal Metropolis-Hastings steps, which re-filter the observed prefix for each proposal.

Parameters:

Name	Description	Default
`model`	A :class:`StateSpaceModel`, :class:`LinearGaussianModel` or :class:`TemporalFilter`.	required
`logprior_fn`	`theta -> scalar` log prior. Initial parameters passed to `init` must be draws from this prior.	required
`num_particles`	Latent particles per filter (particle filters only).	`None`
`ess_threshold`	Resample when the ESS falls below this fraction.	`0.5`
`num_moves`	MH moves after each resampling; `0` disables them.	`1`
`proposal_fn`	`(key, theta) -> (theta', log_q_reverse_minus_forward)`. Defaults to a Gaussian random walk with the weighted population covariance and a scale tuned towards 23.4% acceptance.	`None`
`batch_size`	Parameters processed at once; `0` vectorizes all.	`0`

Source code in probjax/inference/smc/temporal.py

def temporal_smc(
    model,
    logprior_fn,
    *,
    num_particles=None,
    ess_threshold=0.5,
    num_moves=1,
    proposal_fn=None,
    batch_size=0,
):
    """SMC over parameters that assimilates one observation per step.

    Each parameter particle carries its own filter: a particle filter for a
    :class:`StateSpaceModel` (SMC², with ``num_particles`` latent particles) or
    a Kalman filter for a :class:`LinearGaussianModel` (exact). After every
    observation the outer weights are multiplied by the filters' likelihood
    increments. When the ESS drops below ``ess_threshold``, parameters are
    resampled together with their filters and moved by ``num_moves``
    pseudo-marginal Metropolis-Hastings steps, which re-filter the observed
    prefix for each proposal.

    Args:
        model: A :class:`StateSpaceModel`, :class:`LinearGaussianModel` or
            :class:`TemporalFilter`.
        logprior_fn: ``theta -> scalar`` log prior. Initial parameters passed to
            ``init`` must be draws from this prior.
        num_particles: Latent particles per filter (particle filters only).
        ess_threshold: Resample when the ESS falls below this fraction.
        num_moves: MH moves after each resampling; ``0`` disables them.
        proposal_fn: ``(key, theta) -> (theta', log_q_reverse_minus_forward)``.
            Defaults to a Gaussian random walk with the weighted population
            covariance and a scale tuned towards 23.4% acceptance.
        batch_size: Parameters processed at once; ``0`` vectorizes all.
    """
    filt = temporal_filter(model, num_particles)
    if filt.likelihood_kind not in ("exact", "unbiased"):
        raise ValueError("temporal_smc needs an exact or unbiased likelihood.")
    if not 0 <= ess_threshold <= 1:
        raise ValueError("ess_threshold must lie in [0, 1].")
    ensure_int("num_moves", num_moves, min_value=0)
    adaptive = proposal_fn is None

    def init(key, parameters, t0=0.0):
        leaves = jax.tree.leaves(parameters)
        if not leaves or any(x.ndim < 1 for x in leaves):
            raise ValueError("Parameter leaves must have a leading particle axis.")
        n = leaves[0].shape[0]
        if n < 1 or any(x.shape[0] != n for x in leaves):
            raise ValueError("All parameter leaves must have the same nonempty batch.")
        t0 = jnp.asarray(t0, dtype=jnp.result_type(t0, 0.0))
        filter_states = map_fn(lambda a: filt.init(a[0], a[1], t0), batch_size)((
            jax.random.split(key, n),
            parameters,
        ))
        log_priors = map_fn(logprior_fn, batch_size)(parameters)
        if log_priors.shape != (n,):
            raise ValueError("logprior_fn must return a scalar per parameter.")
        zeros = jnp.zeros_like(log_priors)
        return TemporalSMCState(
            parameters,
            filter_states,
            zeros - jnp.log(n),
            zeros,
            jnp.zeros((), zeros.dtype),
            t0,
            t0,
            jnp.array(0, jnp.int32),
            jnp.ones((), zeros.dtype),
        )

    def rejuvenate(key, state, parameters, filter_states, log_likelihoods, data):
        count = state.num_observations + 1  # includes the new observation
        if adaptive:
            covariance = population_covariance(state.parameters, state.log_weights)
            factor = jnp.linalg.cholesky(covariance)

        def propose(key, theta):
            if not adaptive:
                return proposal_fn(key, theta)
            flat, unravel = ravel_pytree(theta)
            noise = factor @ jax.random.normal(key, flat.shape)
            scale = state.proposal_scale * _SCALE_FACTOR / jnp.sqrt(flat.size)
            return unravel(flat + scale * noise), 0.0

        def move(_, carry):
            key, theta, filter_state, log_likelihood, accepted = carry
            key, proposal_key, replay_key, accept_key = jax.random.split(key, 4)
            proposed, correction = propose(proposal_key, theta)
            proposed_state, proposed_ll = replay_filter(
                filt, replay_key, proposed, state.t0, data, count
            )
            log_ratio = (
                logprior_fn(proposed)
                - logprior_fn(theta)
                + proposed_ll
                - log_likelihood
                + correction
            )
            # The current estimate is kept on rejection: refreshing it would
            # change the pseudo-marginal target.
            current, (accept, _, _) = static_binomial_sampling(
                accept_key,
                jnp.where(jnp.isnan(log_ratio), -jnp.inf, log_ratio),
                (theta, filter_state, log_likelihood),
                (proposed, proposed_state, proposed_ll),
            )
            return (key, *current, accepted + accept)

        def moves(args):
            key, theta, filter_state, log_likelihood = args
            carry = (key, theta, filter_state, log_likelihood, jnp.array(0.0))
            _, theta, filter_state, log_likelihood, accepted = jax.lax.fori_loop(
                0, num_moves, move, carry
            )
            return theta, filter_state, log_likelihood, accepted / num_moves

        n = state.log_weights.shape[0]
        parameters, filter_states, log_likelihoods, rates = map_fn(moves, batch_size)((
            jax.random.split(key, n),
            parameters,
            filter_states,
            log_likelihoods,
        ))
        return parameters, filter_states, log_likelihoods, rates.mean()

    def step(key, state, ts, observations, observed=None):
        """Assimilate observation ``state.num_observations`` of the full data.

        Pass the whole dataset seen so far: the new observation is read at
        index ``state.num_observations`` and moves replay the prefix before it.
        """
        data = validate_observations(ts, observations, observed)
        ts, observations, observed = data
        count = state.num_observations
        n = state.log_weights.shape[0]
        index = jnp.minimum(count, ts.shape[0] - 1)
        t = ts[index]
        filter_key, resample_key, move_key = jax.random.split(key, 3)

        filter_states, infos = map_fn(
            lambda a: filt.step(
                a[0], a[1], a[2], t, observations[index], observed[index]
            ),
            batch_size,
        )((
            jax.random.split(filter_key, n),
            state.filter_states,
            state.parameters,
        ))
        # A NaN increment (e.g. a simulator diverging at extreme parameters)
        # gives that particle zero weight instead of poisoning the population.
        increments = jnp.where(
            jnp.isnan(infos.log_likelihood), -jnp.inf, infos.log_likelihood
        )
        log_likelihoods = state.log_likelihoods + increments
        log_weights = state.log_weights + increments
        increment = jax.scipy.special.logsumexp(log_weights)
        # The data must extend what the state has seen: its last assimilated
        # time is state.t.
        consistent = (count == 0) | (ts[jnp.maximum(count - 1, 0)] == state.t)
        valid = (count < ts.shape[0]) & (t > state.t) & consistent
        valid = valid & jnp.isfinite(increment)
        log_weights = log_weights - increment
        ess = effective_sample_size(log_weights)
        resampled = valid & (ess < ess_threshold * n)

        ancestors = jax.lax.cond(
            resampled,
            lambda: systematic(resample_key, jnp.exp(log_weights), n),
            lambda: jnp.arange(n),
        ).astype(jnp.int32)
        parameters, filter_states, log_likelihoods = jax.tree.map(
            lambda x: x[ancestors], (state.parameters, filter_states, log_likelihoods)
        )
        weighted = state._replace(log_weights=log_weights)
        log_weights = jnp.where(resampled, -jnp.log(n), log_weights)
        acceptance_rate = jnp.zeros((), log_weights.dtype)
        if num_moves > 0:
            parameters, filter_states, log_likelihoods, acceptance_rate = jax.lax.cond(
                resampled,
                lambda: rejuvenate(
                    move_key,
                    weighted,
                    parameters,
                    filter_states,
                    log_likelihoods,
                    data,
                ),
                lambda: (
                    parameters,
                    filter_states,
                    log_likelihoods,
                    acceptance_rate,
                ),
            )
        proposal_scale = state.proposal_scale
        if adaptive and num_moves > 0:
            tuned = update_scale_from_acceptance_rate(
                proposal_scale, acceptance_rate, _TARGET_ACCEPTANCE
            )
            proposal_scale = jnp.where(
                resampled, jnp.clip(tuned, 1e-3, 1e3), proposal_scale
            )

        next_state = TemporalSMCState(
            parameters,
            filter_states,
            log_weights,
            log_likelihoods,
            state.log_evidence + increment,
            state.t0,
            t,
            count + 1,
            proposal_scale,
        )
        # Atomic: an invalid step (impossible observation, nothing left to
        # assimilate) leaves the state unchanged.
        next_state = jax.lax.cond(valid, lambda: next_state, lambda: state)
        info = TemporalSMCInfo(
            jnp.where(valid, increment, 0.0),
            jnp.where(valid, ess, 0.0),
            resampled,
            ancestors,
            acceptance_rate,
            valid,
            jnp.array(False),
        )
        return next_state, info

    return TemporalSMC(init, step)

probjax.inference.run_temporal_smc ¶

run_temporal_smc(kernel, key, state, ts, observations, *, observed=None, history='none')

Assimilate every observation at ts that state has not seen yet.

To continue a run, call again with the longer dataset: entries already assimilated are skipped without consuming random keys, so two calls give the same result as one. history is 'none', 'full' or a window size; result.log_likelihood covers only this call's observations.

Source code in probjax/inference/smc/temporal.py

def run_temporal_smc(
    kernel, key, state, ts, observations, *, observed=None, history="none"
):
    """Assimilate every observation at ``ts`` that ``state`` has not seen yet.

    To continue a run, call again with the longer dataset: entries already
    assimilated are skipped without consuming random keys, so two calls give
    the same result as one. ``history`` is ``'none'``, ``'full'`` or a window
    size; ``result.log_likelihood`` covers only this call's observations.
    """
    data = validate_observations(ts, observations, observed)
    ts = data[0]
    zero = jnp.zeros((), state.log_weights.dtype)
    if ts.shape[0] == 0:
        return TemporalResult(state, zero, key, None)
    n = state.log_weights.shape[0]
    skip_info = TemporalSMCInfo(
        zero,
        zero,
        jnp.array(False),
        jnp.arange(n, dtype=jnp.int32),
        zero,
        jnp.array(True),
        jnp.array(True),
    )

    def step(key, state, i):
        return jax.lax.cond(
            i < state.num_observations,
            lambda: (state, skip_info),
            lambda: kernel.step(key, state, *data),
        )

    return scan_history(step, key, state, ts, jnp.arange(ts.shape[0]), history)

probjax.inference.sample_paths ¶

sample_paths(key, trace, model, theta, *, num_samples=1, method=None)

Draw joint latent paths from a filter trace at fixed theta.

Paths have shape (len(trace.ts) + 1, num_samples, D), including the initial state. A :class:LinearGaussianModel trace is sampled exactly backward. A particle trace uses method='backward' (FFBSi, cost O(TNM), needs model.transition_logdensity) or 'ancestry' (stored parents, O(T*M), prone to path collapse). By default backward sampling is used whenever the model has a transition density.

Source code in probjax/inference/filtering/trajectory.py

def sample_paths(key, trace, model, theta, *, num_samples=1, method=None):
    """Draw joint latent paths from a filter trace at fixed ``theta``.

    Paths have shape ``(len(trace.ts) + 1, num_samples, D)``, including the
    initial state. A :class:`LinearGaussianModel` trace is sampled exactly
    backward. A particle trace uses ``method='backward'`` (FFBSi, cost
    O(T*N*M), needs ``model.transition_logdensity``) or ``'ancestry'``
    (stored parents, O(T*M), prone to path collapse). By default backward
    sampling is used whenever the model has a transition density.
    """
    if num_samples < 1:
        raise ValueError("num_samples must be positive.")
    if isinstance(model, LinearGaussianModel):
        return _sample_gaussian_paths(key, trace, model, theta, num_samples)
    if method is None:
        method = "ancestry" if model.transition_logdensity is None else "backward"
    states = _prepend(trace.initial_state, trace.states)
    if method == "ancestry":
        return _genealogies(
            key,
            states.particles,
            states.log_weights[-1],
            trace.infos.ancestors,
            num_samples,
        )
    if method != "backward" or model.transition_logdensity is None:
        raise ValueError("Choose 'ancestry', or 'backward' with a transition density.")
    times = jnp.concatenate((jnp.asarray(trace.initial_state.t)[None], trace.ts))
    paths, _ = particle_smoother(
        key,
        times,
        states.particles,
        states.log_weights,
        lambda new, old, a, b: model.transition_logdensity(theta, new, old, a, b),
        num_samples=num_samples,
    )
    return paths

probjax.inference.smooth_gaussian_path ¶

smooth_gaussian_path(trace, model, theta)

RTS smoothed means and covariances of a :class:LinearGaussianModel trace.

Both include the initial state, so they have len(trace.ts) + 1 entries.

Source code in probjax/inference/filtering/trajectory.py

def smooth_gaussian_path(trace, model, theta):
    """RTS smoothed means and covariances of a :class:`LinearGaussianModel` trace.

    Both include the initial state, so they have ``len(trace.ts) + 1`` entries.
    """
    states, gains = _gaussian_backward_data(trace, model, theta)

    def backward(carry, data):
        next_mean, next_cov = carry
        mean, cov, pred_mean, pred_cov, gain = data
        mean = mean + gain @ (next_mean - pred_mean)
        cov = cov + gain @ (next_cov - pred_cov) @ gain.T
        cov = (cov + cov.T) / 2
        return (mean, cov), (mean, cov)

    _, (means, covs) = jax.lax.scan(
        backward,
        (states.mean[-1], states.cov[-1]),
        (
            states.mean[:-1],
            states.cov[:-1],
            trace.infos.mean_pred,
            trace.infos.cov_pred,
            gains,
        ),
        reverse=True,
    )
    return (
        jnp.concatenate((means, states.mean[-1:])),
        jnp.concatenate((covs, states.cov[-1:])),
    )

probjax.inference.sample_joint_paths ¶

sample_joint_paths(key, state, model, ts, observations, *, observed=None, num_samples=100, num_particles=None, method=None, batch_size=0)

Draw parameters from a temporal SMC state, then a path for each.

Pass the same data as to :func:run_temporal_smc; only the observations the state has assimilated are used. Paths have shape (num_samples, num_assimilated + 1, D). Gaussian paths are exact given each drawn parameter; particle paths re-run a filter with num_particles and so add a finite-particle approximation. Must be called outside jit, since the assimilated count sets array shapes.

Source code in probjax/inference/filtering/joint.py

def sample_joint_paths(
    key,
    state,
    model,
    ts,
    observations,
    *,
    observed=None,
    num_samples=100,
    num_particles=None,
    method=None,
    batch_size=0,
):
    """Draw parameters from a temporal SMC state, then a path for each.

    Pass the same data as to :func:`run_temporal_smc`; only the observations
    the state has assimilated are used. Paths have shape
    ``(num_samples, num_assimilated + 1, D)``. Gaussian paths are exact given
    each drawn parameter; particle paths re-run a filter with
    ``num_particles`` and so add a finite-particle approximation. Must be
    called outside ``jit``, since the assimilated count sets array shapes.
    """
    if num_samples < 1:
        raise ValueError("num_samples must be positive.")
    if is_tracing(state.num_observations):
        raise ValueError("sample_joint_paths must be called outside jit.")
    count = int(state.num_observations)
    ts, observations, observed = validate_observations(ts, observations, observed)
    data = (ts[:count], observations[:count], observed[:count])
    select_key, path_key = jax.random.split(key)
    indices = jax.random.categorical(
        select_key, state.log_weights, shape=(num_samples,)
    )
    parameters = jax.tree.map(lambda x: x[indices], state.parameters)

    def draw(args):
        key, theta = args
        filter_key, sample_key = jax.random.split(key)
        result = run_temporal_filter(
            filter_key,
            model,
            theta,
            data[0],
            data[1],
            t0=state.t0,
            num_particles=num_particles,
            observed=data[2],
        )
        return sample_paths(sample_key, result.trace, model, theta, method=method)[:, 0]

    paths = map_fn(draw, batch_size)((
        jax.random.split(path_key, num_samples),
        parameters,
    ))
    return JointPathSamples(parameters, paths, indices)

probjax.inference.particle_gibbs ¶

particle_gibbs(key, reference_path, model, theta, ts, observations, *, t0=0.0, num_particles=32, observed=None, ancestor_sampling=True)

One conditional-SMC update (PGAS by default), returning a new joint path.

reference_path has shape (len(ts) + 1, D), including x(t0), and model is a :class:StateSpaceModel with array states. Ancestor sampling needs model.transition_logdensity. Multinomial resampling at every step is intentional: pinning one index in systematic resampling is not a valid conditional resampler.

The update preserves the smoothing law at fixed theta. Repeated calls form an MCMC chain; one call is not an independent posterior draw.

Source code in probjax/inference/filtering/trajectory.py

def particle_gibbs(
    key,
    reference_path,
    model,
    theta,
    ts,
    observations,
    *,
    t0=0.0,
    num_particles=32,
    observed=None,
    ancestor_sampling=True,
):
    """One conditional-SMC update (PGAS by default), returning a new joint path.

    ``reference_path`` has shape ``(len(ts) + 1, D)``, including ``x(t0)``, and
    ``model`` is a :class:`StateSpaceModel` with array states. Ancestor
    sampling needs ``model.transition_logdensity``. Multinomial resampling at
    every step is intentional: pinning one index in systematic resampling is
    not a valid conditional resampler.

    The update preserves the smoothing law at fixed ``theta``. Repeated calls
    form an MCMC chain; one call is not an independent posterior draw.
    """
    ts, observations, observed = validate_observations(ts, observations, observed)
    if num_particles < 2:
        raise ValueError("Conditional SMC requires at least two particles.")
    if reference_path.ndim != 2 or reference_path.shape[0] != ts.shape[0] + 1:
        raise ValueError("reference_path must have shape (len(ts)+1, D).")
    if ancestor_sampling and model.transition_logdensity is None:
        raise ValueError("Ancestor sampling needs model.transition_logdensity.")
    key, initial_key = jax.random.split(key)
    particles = jax.vmap(model.init, (0, None))(
        jax.random.split(initial_key, num_particles), theta
    )
    particles = particles.at[0].set(reference_path[0])
    initial_particles = particles
    weights = jnp.full((num_particles,), -jnp.log(num_particles))
    transition = jax.vmap(model.step, (0, None, 0, None, None))
    likelihood = jax.vmap(model.log_likelihood, (None, 0, None, None))

    def step(carry, data):
        key, particles, weights, previous = carry
        key, resample_key, ancestor_key, predict_key = jax.random.split(key, 4)
        t, y, mask, reference = data
        parents = jax.random.categorical(resample_key, weights, shape=(num_particles,))
        if ancestor_sampling:
            ancestor_weights = weights + jax.vmap(
                lambda x: model.transition_logdensity(theta, reference, x, previous, t)
            )(particles)
            reference_parent = jax.random.categorical(ancestor_key, ancestor_weights)
        else:
            reference_parent = 0
        parents = parents.at[0].set(reference_parent)
        keys = jax.random.split(predict_key, num_particles)
        particles = transition(keys, theta, particles[parents], previous, t)
        particles = particles.at[0].set(reference)
        weights = jax.lax.cond(
            mask,
            lambda: likelihood(theta, particles, y, t),
            lambda: jnp.zeros_like(weights),
        )
        weights = weights - jax.scipy.special.logsumexp(weights)
        return (key, particles, weights, t), (particles, parents)

    (key, _, weights, _), (populations, ancestors) = jax.lax.scan(
        step,
        (key, particles, weights, jnp.asarray(t0, dtype=ts.dtype)),
        (ts, observations, observed, reference_path[1:]),
    )
    populations = jnp.concatenate((initial_particles[None], populations))
    return _genealogies(key, populations, weights, ancestors, 1)[:, 0]

probjax.inference.TemporalFilter ¶

Bases: NamedTuple

Pure init(key, theta, t0) and step(key, state, theta, t, y, observed).

step returns (state, info) with info.log_likelihood the incremental log likelihood. likelihood_kind describes the likelihood, not its logarithm: 'exact' or 'unbiased'.

Source code in probjax/inference/filtering/state_space.py

class TemporalFilter(NamedTuple):
    """Pure ``init(key, theta, t0)`` and ``step(key, state, theta, t, y, observed)``.

    ``step`` returns ``(state, info)`` with ``info.log_likelihood`` the
    incremental log likelihood. ``likelihood_kind`` describes the likelihood,
    not its logarithm: ``'exact'`` or ``'unbiased'``.
    """

    init: Callable
    step: Callable
    likelihood_kind: str

probjax.inference.TemporalTrace ¶

Bases: NamedTuple

Stored states at ts, with their immediately preceding initial state.

Windowed traces condition on the filtering distribution at the window start; they do not provide smoothed estimates for states discarded from the window.

Source code in probjax/inference/filtering/state_space.py

class TemporalTrace(NamedTuple):
    """Stored states at ``ts``, with their immediately preceding initial state.

    Windowed traces condition on the filtering distribution at the window start;
    they do not provide smoothed estimates for states discarded from the window.
    """

    initial_state: Any
    ts: Any
    states: Any
    infos: Any

probjax.inference.TemporalResult ¶

Bases: NamedTuple

Source code in probjax/inference/filtering/state_space.py

class TemporalResult(NamedTuple):
    state: Any
    log_likelihood: Any
    key: Any
    trace: Any

probjax.inference.TemporalSMC ¶

Bases: NamedTuple

Source code in probjax/inference/smc/temporal.py

class TemporalSMC(NamedTuple):
    init: Callable
    step: Callable

probjax.inference.TemporalSMCState ¶

Bases: NamedTuple

Source code in probjax/inference/smc/temporal.py

class TemporalSMCState(NamedTuple):
    parameters: Any
    filter_states: Any
    log_weights: Any
    log_likelihoods: Any
    log_evidence: Any
    t0: Any
    t: Any
    num_observations: Any
    proposal_scale: Any

    def posterior_samples(self, key, num_samples):
        """Draw ``num_samples`` parameters from the weighted population."""
        idx = jax.random.categorical(key, self.log_weights, shape=(num_samples,))
        return jax.tree.map(lambda x: x[idx], self.parameters)

posterior_samples ¶

posterior_samples(key, num_samples)

Draw num_samples parameters from the weighted population.

Source code in probjax/inference/smc/temporal.py

def posterior_samples(self, key, num_samples):
    """Draw ``num_samples`` parameters from the weighted population."""
    idx = jax.random.categorical(key, self.log_weights, shape=(num_samples,))
    return jax.tree.map(lambda x: x[idx], self.parameters)

probjax.inference.TemporalSMCInfo ¶

Bases: NamedTuple

Source code in probjax/inference/smc/temporal.py

class TemporalSMCInfo(NamedTuple):
    log_likelihood: Any
    ess: Any
    resampled: Any
    ancestors: Any
    acceptance_rate: Any
    valid: Any
    skipped: Any

Filtering and smoothing¶

probjax.inference.kalman_filter ¶

Bases: FilterAPI

Kalman filter for a linear Gaussian state space model.

\[dx_t = A_t x_t + B_t dw_t \qquad y_t = \mathcal{N}(y_t; C x_t, R_t)\]

To build a Kalman filter kernel, we require the following components:

Parameters:

Name	Type	Description	Default
`transition_matrix`	`Callable[[float \| ArrayLike], ArrayLike] \| ArrayLike`	Transition matrix A_t	required
`observation_matrix`	`Callable[[float \| ArrayLike], ArrayLike] \| ArrayLike`	Observation matrix C_t	required
`observation_covariance`	`Callable[[float \| ArrayLike], ArrayLike] \| ArrayLike`	Observation covariance matrix R_t	required

Source code in probjax/inference/filtering/kalman_filter.py

class kalman_filter(FilterAPI):
    r"""
    Kalman filter for a linear Gaussian state space model.

    $$dx_t = A_t x_t + B_t dw_t \qquad y_t = \mathcal{N}(y_t; C x_t, R_t)$$

    To build a Kalman filter kernel, we require the following components:

    Args:
        transition_matrix (Callable[[float | ArrayLike], ArrayLike] | ArrayLike):
            Transition matrix A_t
        transition_covariance_matrix (Callable[[float | ArrayLike], ArrayLike] |
            ArrayLike): Transition covariance matrix Q_t
        observation_matrix (Callable[[float | ArrayLike], ArrayLike] | ArrayLike):
            Observation matrix C_t
        observation_covariance (Callable[[float | ArrayLike], ArrayLike] | ArrayLike):
            Observation covariance matrix R_t
    """

    init = init
    build_kernel = build_kernel

    default_unpack = staticmethod(gaussian_unpack)

probjax.inference.extended_kalman_filter ¶

Bases: FilterAPI

Extended Kalman filter for a nonlinear state space model.

\[x_{t+1} = f(x_t, t) + w_t \qquad y_t = \mathcal{N}(y_t; h(x_t, t), R_t)\]

The EKF linearizes the transition and observation functions around the current state estimate to apply the standard Kalman filter update equations.

Parameters:

Name	Type	Description	Default
`transition_model_fn`		(x, cov, t_old, t) -> (x_pred, Phi, Q) Returns the nonlinear predicted state, the Jacobian of the transition, and the process noise covariance.	required
`observation_model_fn`		(x, cov, t) -> (y_pred, C, R) Returns the nonlinear predicted observation, the Jacobian of the observation function, and the observation noise covariance.	required

Helpers for building these callables

make_linearized_transition(f, Q): wraps f(x,t_old,t)->x with auto-Jacobian
make_linearized_observation(h, R): wraps h(x,t)->y with auto-Jacobian
make_continuous_transition(drift, B): continuous SDE via matrix fraction decomposition

Source code in probjax/inference/filtering/extended_kalman_filter.py

class extended_kalman_filter(FilterAPI):
    r"""
    Extended Kalman filter for a nonlinear state space model.

    $$x_{t+1} = f(x_t, t) + w_t \qquad y_t = \mathcal{N}(y_t; h(x_t, t), R_t)$$

    The EKF linearizes the transition and observation functions around the current
    state estimate to apply the standard Kalman filter update equations.

    Args:
        transition_model_fn: (x, cov, t_old, t) -> (x_pred, Phi, Q)
            Returns the nonlinear predicted state, the Jacobian of the
            transition, and the process noise covariance.
        observation_model_fn: (x, cov, t) -> (y_pred, C, R)
            Returns the nonlinear predicted observation, the Jacobian of the
            observation function, and the observation noise covariance.

    Helpers for building these callables:
        - `make_linearized_transition(f, Q)`: wraps f(x,t_old,t)->x with auto-Jacobian
        - `make_linearized_observation(h, R)`: wraps h(x,t)->y with auto-Jacobian
        - `make_continuous_transition(drift, B)`: continuous SDE via matrix
          fraction decomposition
    """

    init = gaussian_init
    build_kernel = build_kernel

    default_unpack = staticmethod(gaussian_unpack)

probjax.inference.filtering.unscented_kalman_filter ¶

ukf ¶

Bases: FilterAPI

Unscented Kalman filter inference algorithm.

This class implements the unscented Kalman filter algorithm. The unscented Kalman filter is a generalization of the Kalman filter to non-linear and non-Gaussian models.

To build an unscented Kalman filter, you need to provide the following functions: Args: transition_fn (Callable): Transition function f(x_t, t) -> x_{t+1} transition_covariance_matrix (Callable | ArrayLike): Transition covariance matrix Q(t) or Q observation_fn (Callable): Observation function h(x_t, t) -> y_t observation_covariance (Callable | ArrayLike): Observation covariance matrix R(t) or R

Source code in probjax/inference/filtering/unscented_kalman_filter.py

class ukf(FilterAPI):
    """Unscented Kalman filter inference algorithm.

    This class implements the unscented Kalman filter algorithm. The unscented Kalman
    filter is a generalization of the Kalman filter to non-linear and non-Gaussian
    models.

    To build an unscented Kalman filter, you need to provide the following functions:
    Args:
        transition_fn (Callable): Transition function f(x_t, t) -> x_{t+1}
        transition_covariance_matrix (Callable | ArrayLike): Transition covariance
            matrix Q(t) or Q
        observation_fn (Callable): Observation function h(x_t, t) -> y_t
        observation_covariance (Callable | ArrayLike): Observation covariance matrix
            R(t) or R
    """

    init = init
    build_kernel = build_kernel

    default_unpack = staticmethod(gaussian_unpack)

merwe_sigma_point ¶

merwe_sigma_point(mu0, cov0, alpha=1.0, beta=2.0, kappa=0.0)

Generate sigma points for unscented Kalman filter, see [1].

Parameters:

Name	Type	Description	Default
`mu0`	`ArrayLike`	Mean of the state	required
`cov0`	`ArrayLike`	Covariance of the state	required
`alpha`	`ArrayLike`	Determines the spread around the mean (small values lead to large weight which require high precission (64 bit)). Large values will spread the sigma points further, typically letting to an overestimation of the covariance, small values will lead to an underestimation. Literature suggests 1e-3 but this requires 64 bit precision to work at all... . Defaults to 1..	`1.0`
`beta`	`ArrayLike`	Prior on covariance. Defaults to 2..	`2.0`
`kappa`	`ArrayLike`	Additional parameter. Defaults to 0..	`0.0`

Returns:

Type	Description
`Tuple[ArrayLike, ArrayLike, ArrayLike]`	Tuple[NDArray, NDArray, NDArray]: Sigma points, weights_mean, weights_cov

References

[1] R. Van der Merwe "Sigma-Point Kalman Filters for Probabilitic Inference in Dynamic State-Space Models" (Doctoral dissertation)

Source code in probjax/inference/filtering/unscented_kalman_filter.py

def merwe_sigma_point(
    mu0: ArrayLike,
    cov0: ArrayLike,
    alpha: ArrayLike = 1.0,
    beta: ArrayLike = 2.0,
    kappa: ArrayLike = 0.0,
) -> Tuple[ArrayLike, ArrayLike, ArrayLike]:
    """Generate sigma points for unscented Kalman filter, see [1].

    Args:
        mu0 (ArrayLike): Mean of the state
        cov0 (ArrayLike): Covariance of the state
        alpha (ArrayLike, optional): Determines the spread around the mean
            (small values lead to large weight which require high precission (64 bit)).
            Large values will spread the sigma points further, typically letting
            to an overestimation of the covariance, small values will lead to an
            underestimation. Literature suggests 1e-3 but this requires 64 bit
            precision to work at all... . Defaults to 1..
        beta (ArrayLike, optional): Prior on covariance. Defaults to 2..
        kappa (ArrayLike, optional): Additional parameter. Defaults to 0..

    Returns:
        Tuple[NDArray, NDArray, NDArray]: Sigma points, weights_mean, weights_cov

    References:
        [1] R. Van der Merwe "Sigma-Point Kalman Filters for Probabilitic
            Inference in Dynamic State-Space Models" (Doctoral dissertation)
    """

    with jax.ensure_compile_time_eval():
        D = mu0.shape[0]
        lambda_ = alpha**2 * (D + kappa) - D
        weights_mean_0 = jnp.atleast_1d(lambda_ / (D + lambda_))
        weights_cov_0 = jnp.atleast_1d(weights_mean_0 + (1 - alpha**2 + beta))
        weights_mean_cov_1_L = 1 / (2 * (D + lambda_)) * jnp.ones(2 * D)

        weights_mean = jnp.concatenate([weights_mean_0, weights_mean_cov_1_L], axis=0)
        weights_cov = jnp.concatenate([weights_cov_0, weights_mean_cov_1_L], axis=0)

    sqrt_cov = jnp.linalg.cholesky((D + lambda_) * cov0)
    sigma_points_0 = mu0[None, :]
    sigma_points_1_L = mu0 + sqrt_cov.T
    sigma_points_L_2L = mu0 - sqrt_cov.T
    sigma_points = jnp.concatenate(
        [sigma_points_0, sigma_points_1_L, sigma_points_L_2L], axis=0
    )

    return sigma_points, weights_mean, weights_cov

julier_uhlmann_sigma_points ¶

julier_uhlmann_sigma_points(mu0, cov0, kappa=0.0)

Generate Julier-Uhlmann sigma points.

Parameters:

Name	Type	Description	Default
`mu0`	`ndarray`	Mean of the state (D, )	required
`cov0`	`ndarray`	Covariance of the state (D, D)	required
`kappa`	`float`	Spread parameter. Often chosen as kappa = 3 - D for state dimension D.	`0.0`

Returns:

Type	Description
`Tuple[ndarray, ndarray, ndarray]`	Tuple[sigma_points, weights_mean, weights_cov]: sigma_points: (2D+1, D) weights_mean: (2D+1,) weights_cov: (2D+1,)

Source code in probjax/inference/filtering/unscented_kalman_filter.py

def julier_uhlmann_sigma_points(
    mu0: jnp.ndarray,
    cov0: jnp.ndarray,
    kappa: float = 0.0,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """Generate Julier-Uhlmann sigma points.

    Args:
        mu0 (jnp.ndarray): Mean of the state (D, )
        cov0 (jnp.ndarray): Covariance of the state (D, D)
        kappa (float): Spread parameter.
            Often chosen as kappa = 3 - D for state dimension D.

    Returns:
        Tuple[sigma_points, weights_mean, weights_cov]:
            sigma_points: (2D+1, D)
            weights_mean: (2D+1,)
            weights_cov: (2D+1,)
    """
    D = mu0.shape[0]
    lambda_ = D + kappa
    # Compute sqrt of scaled covariance
    sqrt_cov = jnp.linalg.cholesky(lambda_ * cov0)

    # Sigma points
    sigma_points_0 = mu0[None, :]
    sigma_points_pos = mu0 + sqrt_cov.T
    sigma_points_neg = mu0 - sqrt_cov.T
    sigma_points = jnp.concatenate(
        [sigma_points_0, sigma_points_pos, sigma_points_neg], axis=0
    )

    # Weights
    w0 = kappa / (D + kappa)
    wi = 1.0 / (2.0 * (D + kappa)) * jnp.ones(2 * D)
    weights_mean = jnp.concatenate([jnp.array([w0]), wi], axis=0)
    weights_cov = jnp.concatenate([jnp.array([w0]), wi], axis=0)

    return sigma_points, weights_mean, weights_cov

spherical_simplex_sigma_points ¶

spherical_simplex_sigma_points(mu0, cov0)

Generate spherical simplex sigma points without explicit Python loops.

This creates a (D+1) x D array of points that form a regular simplex. The steps are: 1. Construct an initial (D+1, D) array corresponding to a simplex. 2. Center the points so they have zero mean. 3. Scale the points to achieve unit covariance. 4. Apply the square root of the covariance (via Cholesky) and then add mu0.

Parameters:

Name	Type	Description	Default
`mu0`	`ndarray`	Mean of the state, shape (D,)	required
`cov0`	`ndarray`	Covariance of the state, shape (D,D)	required

Returns:

Name	Type	Description
`sigma_points`	`ndarray`	(D+1, D)
`weights_mean`	`ndarray`	(D+1,)
`weights_cov`	`ndarray`	(D+1,)

Source code in probjax/inference/filtering/unscented_kalman_filter.py

def spherical_simplex_sigma_points(
    mu0: jnp.ndarray, cov0: jnp.ndarray
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """
    Generate spherical simplex sigma points without explicit Python loops.

    This creates a (D+1) x D array of points that form a regular simplex.
    The steps are:
    1. Construct an initial (D+1, D) array corresponding to a simplex.
    2. Center the points so they have zero mean.
    3. Scale the points to achieve unit covariance.
    4. Apply the square root of the covariance (via Cholesky) and then add mu0.

    Args:
        mu0 (jnp.ndarray): Mean of the state, shape (D,)
        cov0 (jnp.ndarray): Covariance of the state, shape (D,D)

    Returns:
        sigma_points (jnp.ndarray): (D+1, D)
        weights_mean (jnp.ndarray): (D+1,)
        weights_cov (jnp.ndarray): (D+1,)
    """
    D = mu0.shape[0]
    # A Helmert basis spans the subspace orthogonal to the all-ones
    # vector. Its columns are orthonormal, so equal weights recover I.
    rows = jnp.arange(D + 1)[:, None]
    cols = jnp.arange(D)[None, :]
    denominator = jnp.sqrt((cols + 1) * (cols + 2))
    basis = jnp.where(rows <= cols, 1.0, jnp.where(rows == cols + 1, -(cols + 1), 0.0))
    points = jnp.sqrt(D + 1) * basis / denominator
    sigma_points = mu0[None, :] + points @ jnp.linalg.cholesky(cov0).T

    # Equal weights for mean and covariance
    weights_mean = jnp.ones(D + 1) / (D + 1)
    weights_cov = jnp.ones(D + 1) / (D + 1)

    return sigma_points, weights_mean, weights_cov

unscented_transform ¶

unscented_transform(sigma_points, weights_mean, weights_cov, noise_cov=None, mean_fn=None, cov_fn=None)

Unscented transform

Parameters:

Name	Type	Description	Default
`sigma_points`	`ArrayLike`	(Transformed) sigma points (2D+1, D)	required
`weights_mean`	`ArrayLike`	Weights for the mean (2D+1,)	required
`weights_cov`	`ArrayLike`	Weights for the covariance (2D+1,)	required
`noise_cov`	`Optional[ArrayLike]`	Additive noise covariance (D,D). Defaults to None.	`None`
`mean_fn`	`Optional[Callable]`	Custom mean_fn predictor mean_fn(simga_point, weight_mean). Defaults to None.	`None`
`cov_fn`	`Optional[Callable]`	Custom cov_fn predict cov_fn(sigma_points, mean, weights_cov). Defaults to None.	`None`

Returns:

Type	Description
`Tuple[ArrayLike, ArrayLike]`	Tuple[NDArray, NDArray]: description

Source code in probjax/inference/filtering/unscented_kalman_filter.py

def unscented_transform(
    sigma_points: ArrayLike,
    weights_mean: ArrayLike,
    weights_cov: ArrayLike,
    noise_cov: Optional[ArrayLike] = None,
    mean_fn: Optional[Callable] = None,
    cov_fn: Optional[Callable] = None,
) -> Tuple[ArrayLike, ArrayLike]:
    """Unscented transform

    Args:
        sigma_points (ArrayLike): (Transformed) sigma points (2D+1, D)
        weights_mean (ArrayLike): Weights for the mean (2D+1,)
        weights_cov (ArrayLike): Weights for the covariance (2D+1,)
        noise_cov (Optional[ArrayLike], optional): Additive noise covariance (D,D).
            Defaults to None.
        mean_fn (Optional[Callable], optional): Custom mean_fn predictor
            mean_fn(simga_point, weight_mean). Defaults to None.
        cov_fn (Optional[Callable], optional): Custom cov_fn predict
            cov_fn(sigma_points, mean, weights_cov). Defaults to None.

    Returns:
        Tuple[NDArray, NDArray]: _description_
    """
    if mean_fn is not None:
        mean = mean_fn(sigma_points, weights_mean)
    else:
        mean = jnp.dot(weights_mean, sigma_points)

    if cov_fn is not None:
        cov = cov_fn(sigma_points, mean, weights_cov)
    else:
        diff = sigma_points - mean[None, :]
        cov = jnp.dot(weights_cov * diff.T, diff)

    if noise_cov is not None:
        cov += noise_cov

    return mean, cov

init ¶

init(mu0, cov0, t=None)

Initialize the unscented Kalman filter.

Parameters:

Name	Type	Description	Default
`mu0`	`ArrayLike`	Initial mean of the state	required
`cov0`	`ArrayLike`	Initial covariance of the state	required
`t`	`Optional[float \| int]`	Time. Defaults to None.	`None`

Returns:

Name	Type	Description
`UnscentedKalmanFilterState`	`UnscentedKalmanFilterState`	Initial state of the unscented Kalman filter

Source code in probjax/inference/filtering/unscented_kalman_filter.py

def init(
    mu0: ArrayLike, cov0: ArrayLike, t: Optional[float | int] = None
) -> UnscentedKalmanFilterState:
    """Initialize the unscented Kalman filter.

    Args:
        mu0 (ArrayLike): Initial mean of the state
        cov0 (ArrayLike): Initial covariance of the state
        t (Optional[float | int], optional): Time. Defaults to None.

    Returns:
        UnscentedKalmanFilterState: Initial state of the unscented Kalman filter
    """
    return UnscentedKalmanFilterState(mu0, cov0, t)

build_kernel ¶

build_kernel(transition_fn, transition_covariance_matrix, observation_fn, observation_covariance, sigma_point_fn=merwe_sigma_point)

Build an unscented Kalman filter kernel.

Parameters:

Name	Type	Description	Default
`transition_fn`	`Callable`	General transition function f(x_t, t, t+1) -> x_{t+1}	required
`transition_covariance_matrix`	`Callable \| ArrayLike`	Transition covariance	required
`matrix Q`	`t, t+1) or Q observation_fn (Callable`	General observation function	required
`h`	`x_t, t) -> y_t observation_covariance (Callable \| ArrayLike`	Observation	required
`sigma_point_fn`	`Callable`	How to generate sigma points. Defaults to	`merwe_sigma_point`

Returns:

Name	Type	Description
`Callable`	`Callable`	Unscented Kalman filter step

Source code in probjax/inference/filtering/unscented_kalman_filter.py

def build_kernel(
    transition_fn: Callable,
    transition_covariance_matrix: Callable | ArrayLike,
    observation_fn: Callable,
    observation_covariance: Callable | ArrayLike,
    sigma_point_fn: Callable = merwe_sigma_point,
) -> Callable:
    """Build an unscented Kalman filter kernel.

    Args:
        transition_fn (Callable): General transition function f(x_t, t, t+1) -> x_{t+1}
        transition_covariance_matrix (Callable | ArrayLike): Transition covariance
        matrix Q(t, t+1) or Q observation_fn (Callable): General observation function
        h(x_t, t) -> y_t  observation_covariance (Callable | ArrayLike): Observation
        covariance matrix R(t) or R
        sigma_point_fn (Callable, optional): How to generate sigma points. Defaults to
        merwe_sigma_point.

    Returns:
        Callable: Unscented Kalman filter step
    """

    def kernel(
        state: UnscentedKalmanFilterState,
        t: Optional[float | int] = None,
        observed: Optional[ArrayLike] = None,
        rng_key: Optional[jnp.ndarray] = None,
    ) -> Tuple[UnscentedKalmanFilterState, UnscentedKalmanFilterInfo]:
        """One step of the unscented Kalman filter.

        Args:
            state (UnscentedKalmanFilterState): Mean and covariance of the state
            t (Optional[float  |  int], optional): Time. Defaults to None.
            observed (Optional[ArrayLike], optional): Observation. Defaults to None.
            rng_key (Optional[jnp.ndarray], optional): Random generator key.
                Defaults to None.

        Returns:
            Tuple[UnscentedKalmanFilterState, UnscentedKalmanFilterInfo]: _description_
        """

        mu0 = state.mean
        cov0 = state.cov
        t_old = state.t
        is_observed = observed is not None

        sigma_points, weights_mean, weights_cov = sigma_point_fn(mu0, cov0)
        predicted_sigma_points = jax.vmap(transition_fn, in_axes=(0, None, None))(
            sigma_points, t_old, t
        )
        Q = call_if_callable(transition_covariance_matrix, t_old, t)

        # Prediction step
        mu1_, cov1_ = unscented_transform(
            predicted_sigma_points, weights_mean, weights_cov, noise_cov=Q
        )

        if is_observed:
            # Include process noise in the observation and cross-covariance.
            predicted_sigma_points, weights_mean, weights_cov = sigma_point_fn(
                mu1_, cov1_
            )
            R = call_if_callable(observation_covariance, t)
            # Observed steps
            y_sigma_points = jax.vmap(observation_fn, in_axes=(0, None))(
                predicted_sigma_points, t
            )
            mu_y, cov_y = unscented_transform(
                y_sigma_points,
                weights_mean,
                weights_cov,
                noise_cov=R,
            )

            # Compute the cross-covariance
            r_state = predicted_sigma_points - mu1_[None, :]
            r_obs = y_sigma_points - mu_y[None, :]
            cross_cov = jnp.dot(weights_cov * r_state.T, r_obs)

            # Compute the Kalman gain
            K = jnp.linalg.solve(cov_y.T, cross_cov.T).T

            # Compute the updated mean and covariance
            r = observed - mu_y
            mu1 = mu1_ + jnp.dot(K, r)
            cov1 = cov1_ - jnp.dot(K, jnp.dot(cov_y, K.T))

            # Compute the log-likelihood
            log_likelihood = innovation_loglik(
                r,
                cov_y,
                solve_fn=lambda S, res: jnp.linalg.solve(S, res),
                logdet_fn=lambda S: jnp.linalg.slogdet(S)[1],
            )

            return UnscentedKalmanFilterState(mu1, cov1, t), UnscentedKalmanFilterInfo(
                mu1_, cov1_, log_likelihood
            )
        else:
            log_likelihood = jnp.array(0.0)
            return UnscentedKalmanFilterState(
                mu1_, cov1_, t
            ), UnscentedKalmanFilterInfo(mu1_, cov1_, log_likelihood)

    return kernel

probjax.inference.filtering.square_root_kf ¶

sq_kalman_filter ¶

Bases: FilterAPI

Square root Kalman filter for a linear Gaussian state space model.

\[dx_t = A_t x_t + B_t dw_t \qquad y_t = \mathcal{N}(y_t; C x_t, R_t)\]

To build a Kalman filter kernel, we require the following components:

Parameters:

Name	Type	Description	Default
`transition_matrix`	`Callable[[float \| ArrayLike], ArrayLike] \| ArrayLike`	Transition matrix A_t	required

Source code in probjax/inference/filtering/square_root_kf.py

class sq_kalman_filter(FilterAPI):
    r"""
    Square root Kalman filter for a linear Gaussian state space model.

    $$dx_t = A_t x_t + B_t dw_t \qquad y_t = \mathcal{N}(y_t; C x_t, R_t)$$

    To build a Kalman filter kernel, we require the following components:

    Args:
        transition_matrix (Callable[[float | ArrayLike], ArrayLike] | ArrayLike):
            Transition matrix A_t
        transition_covariance_matrix_sqrt
            (Callable[[float | ArrayLike], ArrayLike] | ArrayLike): Transition
                covariance matrix Q_t
        observation_matrix
            (Callable[[float | ArrayLike], ArrayLike] | ArrayLike): Observation matrix
            C_t
        observation_covariance
            (Callable[[float | ArrayLike], ArrayLike] | ArrayLike): Observation
            covariance matrix R_t
    """

    init = init
    build_kernel = build_kernel

    @staticmethod
    def default_unpack(state, info):
        return (state.mean, state.std)

sqrt_kf_predict ¶

sqrt_kf_predict(m, CL, A, QL)

Predict step of the square root Kalman filter.

Source code in probjax/inference/filtering/square_root_kf.py

def sqrt_kf_predict(
    m: ArrayLike,
    CL: ArrayLike,
    A: ArrayLike,
    QL: ArrayLike,
) -> Tuple[jax.Array, jax.Array]:
    """
    Predict step of the square root Kalman filter.
    """
    m_new = A @ m
    block = jnp.vstack((A @ CL, QL))
    _, R = jnp.linalg.qr(block, mode="reduced")
    CL_new = R.T
    return m_new, CL_new

sqrt_kf_correct ¶

sqrt_kf_correct(m, CL, H, RL, y, compute_likelihood=False)

Correction step of the square root Kalman filter.

Source code in probjax/inference/filtering/square_root_kf.py

def sqrt_kf_correct(
    m: ArrayLike,
    CL: ArrayLike,
    H: ArrayLike,
    RL: ArrayLike,
    y: ArrayLike,
    compute_likelihood: bool = False,
) -> Tuple[ArrayLike, ArrayLike, float]:
    """
    Correction step of the square root Kalman filter.
    """
    d, D = H.shape

    y_hat = H @ m

    # QR decomposition
    RL_padded = jnp.hstack((RL, jnp.zeros((d, D - d))))

    # Construct QR decomposition block
    block = jnp.vstack((RL_padded, H @ CL))
    _, R = jnp.linalg.qr(block, mode="reduced")

    # Extract submatrices
    SL = jnp.tril(R[:d, :d].T)  # Observation covariance sqrt
    CL_new_factor = jnp.tril(R[d:, d:].T)  # Updated state covariance sqrt

    # Make SL a valid Cholesky factor
    signs = jnp.sign(jnp.diag(SL))
    SL = SL * signs[:, None]
    SL_chol = jax.scipy.linalg.cholesky(SL @ SL.T, lower=True)

    residual = y - y_hat
    Sinv_residual = jnp.linalg.solve(SL, residual)

    m_new = m + jnp.dot(R[:d, d:].T, Sinv_residual)
    CL_new = CL_new_factor

    if compute_likelihood:
        log_likelihood = -0.5 * (
            jnp.dot(residual.T, jnp.linalg.solve(SL_chol, residual))
            + jnp.log(jnp.linalg.det(SL_chol))
        )
    else:
        log_likelihood = jnp.array(0.0)

    return m_new, CL_new, log_likelihood

probjax.inference.rank_reduced_kalman_filter ¶

Bases: FilterAPI

Rank-reduced Kalman filter.

Stores covariance in low-rank form P ~= U S U^T and truncates to a fixed rank.

Source code in probjax/inference/filtering/rank_reduced_kalman_filter.py

class rank_reduced_kalman_filter(FilterAPI):
    """Rank-reduced Kalman filter.

    Stores covariance in low-rank form P ~= U S U^T and truncates to a fixed rank.
    """

    init = init
    build_kernel = build_kernel

    @staticmethod
    def default_unpack(state, info):
        return (state.mean, _factor_to_cov(state.cov_factor, state.cov_core))

probjax.inference.sq_kalman_filter ¶

Bases: FilterAPI

Square root Kalman filter for a linear Gaussian state space model.

\[dx_t = A_t x_t + B_t dw_t \qquad y_t = \mathcal{N}(y_t; C x_t, R_t)\]

To build a Kalman filter kernel, we require the following components:

Parameters:

Name	Type	Description	Default
`transition_matrix`	`Callable[[float \| ArrayLike], ArrayLike] \| ArrayLike`	Transition matrix A_t	required

Source code in probjax/inference/filtering/square_root_kf.py

class sq_kalman_filter(FilterAPI):
    r"""
    Square root Kalman filter for a linear Gaussian state space model.

    $$dx_t = A_t x_t + B_t dw_t \qquad y_t = \mathcal{N}(y_t; C x_t, R_t)$$

    To build a Kalman filter kernel, we require the following components:

    Args:
        transition_matrix (Callable[[float | ArrayLike], ArrayLike] | ArrayLike):
            Transition matrix A_t
        transition_covariance_matrix_sqrt
            (Callable[[float | ArrayLike], ArrayLike] | ArrayLike): Transition
                covariance matrix Q_t
        observation_matrix
            (Callable[[float | ArrayLike], ArrayLike] | ArrayLike): Observation matrix
            C_t
        observation_covariance
            (Callable[[float | ArrayLike], ArrayLike] | ArrayLike): Observation
            covariance matrix R_t
    """

    init = init
    build_kernel = build_kernel

    @staticmethod
    def default_unpack(state, info):
        return (state.mean, state.std)

probjax.inference.ParticleFilter ¶

Bases: FilterAPI

Particle filter inference algorithm.

This class implements the particle filter algorithm. The particle filter is a sequential Monte Carlo method that approximates the filtering distribution of a state-space model. The particle filter is a generalization of the Kalman filter to non-linear and non-Gaussian models.

To build a particle filter, you need to provide the following functions: Args: log_likelihood_fn (Callable): Log likelihood function of the model p(y_t|x_t). transition_fn (Callable): Transition function p(x_t|x_{t-1}) of the model. transition_logdensity_fn (Optional[Callable]): Computes logdensity function of the transition function. Defaults to None. proposal_transition_fn (Optional[Callable]): Transition based on a proposal. Defaults to None. proposal_logdensity_fn (Optional[Callable]): Proposal density function. Defaults to None. resample_criterion (Callable): Criterion to decide when to resample. Defaults to resample_when_ess_below. resample_fn (Callable): Resampling function. Defaults to resample_systematic. unbiased_gradients (bool): Whether to use unbiased gradients. Defaults to False.

Source code in probjax/inference/filtering/particle_filter.py

class ParticleFilter(FilterAPI):
    """Particle filter inference algorithm.

    This class implements the particle filter algorithm. The particle filter is a
    sequential Monte Carlo method that approximates the filtering distribution of a
    state-space model. The particle filter is a generalization of the Kalman filter
    to non-linear and non-Gaussian models.

    To build a particle filter, you need to provide the following functions:
    Args:
        log_likelihood_fn (Callable): Log likelihood function of the model p(y_t|x_t).
        transition_fn (Callable): Transition function p(x_t|x_{t-1}) of the model.
        transition_logdensity_fn (Optional[Callable]): Computes logdensity function of
            the transition function. Defaults to None.
        proposal_transition_fn (Optional[Callable]): Transition based on a proposal.
            Defaults to None.
        proposal_logdensity_fn (Optional[Callable]): Proposal density function.
            Defaults to None.
        resample_criterion (Callable): Criterion to decide when to resample. Defaults
            to resample_when_ess_below.
        resample_fn (Callable): Resampling function. Defaults to resample_systematic.
        unbiased_gradients (bool): Whether to use unbiased gradients. Defaults to
            False.
    """

    init = init
    build_kernel = build_kernel

    @staticmethod
    def default_unpack(state, info):
        return state.particles

probjax.inference.particle_smoother ¶

particle_smoother(key, ts, filter_particles, filter_log_weights, transition_logdensity_fn, ancestors=None, *, num_samples=None)

Forward Filter-Backward Simulator (FFBSi) particle smoother.

Takes the output of a forward particle filter pass and runs a backward simulation pass to approximate the smoothing distribution p(x_{0:T} | y_{1:T}).

The algorithm is O(T * N^2) where T is the number of time steps and N is the number of particles, due to the pairwise transition density evaluation at each backward step.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Random key.	required
`ts`	`Array`	Time grid, shape (T,).	required
`filter_particles`	`Array`	Particles from the filter, shape (T, N, D).	required
`filter_log_weights`	`Array`	Log-weights from the filter, shape (T, N).	required
`transition_logdensity_fn`	`Callable`	Log-density of the transition model. Signature: (x_tp1, x_t, t, tp1) -> scalar log-density.	required
`ancestors`	`Optional[Array]`	Ancestor indices from the filter, shape (T, N). Not used in FFBSi but accepted for API compatibility.	`None`
`num_samples`	`Optional[int]`	Number M of joint trajectories to sample. Defaults to N.	`None`

Returns:

Type	Description
`Tuple[Array, Array]`	Tuple[Array, Array]: - smoothed_particles: shape (T, M, D) - smoothed_log_weights: shape (T, M), uniform weights (1/M)

Source code in probjax/inference/filtering/smoothing.py

def particle_smoother(
    key: Array,
    ts: Array,
    filter_particles: Array,
    filter_log_weights: Array,
    transition_logdensity_fn: Callable,
    ancestors: Optional[Array] = None,
    *,
    num_samples: Optional[int] = None,
) -> Tuple[Array, Array]:
    """Forward Filter-Backward Simulator (FFBSi) particle smoother.

    Takes the output of a forward particle filter pass and runs a backward
    simulation pass to approximate the smoothing distribution p(x_{0:T} | y_{1:T}).

    The algorithm is O(T * N^2) where T is the number of time steps and N is
    the number of particles, due to the pairwise transition density evaluation
    at each backward step.

    Args:
        key (Array): Random key.
        ts (Array): Time grid, shape (T,).
        filter_particles (Array): Particles from the filter, shape (T, N, D).
        filter_log_weights (Array): Log-weights from the filter, shape (T, N).
        transition_logdensity_fn (Callable): Log-density of the transition model.
            Signature: (x_tp1, x_t, t, tp1) -> scalar log-density.
        ancestors (Optional[Array]): Ancestor indices from the filter,
            shape (T, N). Not used in FFBSi but accepted for API compatibility.
        num_samples: Number M of joint trajectories to sample. Defaults to N.

    Returns:
        Tuple[Array, Array]:
            - smoothed_particles: shape (T, M, D)
            - smoothed_log_weights: shape (T, M), uniform weights (1/M)
    """
    T, N, D = filter_particles.shape
    num_samples = N if num_samples is None else num_samples
    if T < 1 or num_samples < 1:
        raise ValueError(
            "A smoother requires a nonempty trace and positive num_samples."
        )

    # Initialize: at the last time step, smoothed = filter
    final_log_weights = filter_log_weights[-1]
    # Sample initial smoothed indices from the final filter distribution
    key, subkey = jax.random.split(key)
    smoothed_indices_T = jax.random.categorical(
        subkey, final_log_weights, shape=(num_samples,)
    )

    # Backward scan from T-1 down to 0
    def backward_step(carry, data):
        key, smoothed_indices_tp1 = carry
        filter_particles_t, filter_log_weights_t, filter_particles_tp1, t, tp1 = data
        key, subkey = jax.random.split(key)

        new_indices = _ffbsi_backward_step(
            subkey,
            smoothed_indices_tp1,
            filter_particles_t,
            filter_log_weights_t,
            filter_particles_tp1,
            transition_logdensity_fn,
            t,
            tp1,
        )
        return (key, new_indices), new_indices

    # Prepare backward scan data: from t=T-2 down to t=0
    # At each step we need: filter_particles[t], filter_log_weights[t],
    #   filter_particles[t+1], ts[t], ts[t+1]
    scan_data = (
        jnp.flip(filter_particles[:-1], axis=0),  # filter particles at t
        jnp.flip(filter_log_weights[:-1], axis=0),  # filter log-weights at t
        jnp.flip(filter_particles[1:], axis=0),  # filter particles at t+1
        jnp.flip(ts[:-1]),  # t
        jnp.flip(ts[1:]),  # t+1
    )

    init_carry = (key, smoothed_indices_T)
    _, all_smoothed_indices = jax.lax.scan(backward_step, init_carry, scan_data)

    # all_smoothed_indices is shape (T-1, N) in reversed time order
    # Flip back to forward time order and append the final indices
    all_smoothed_indices = jnp.flip(all_smoothed_indices, axis=0)
    all_smoothed_indices = jnp.concatenate(
        [all_smoothed_indices, smoothed_indices_T[None, :]], axis=0
    )

    # Gather smoothed particles
    smoothed_particles = jax.vmap(lambda particles, idx: particles[idx])(
        filter_particles, all_smoothed_indices
    )

    # Smoothed weights are uniform
    smoothed_log_weights = jnp.full((T, num_samples), fill_value=-jnp.log(num_samples))

    return smoothed_particles, smoothed_log_weights

probjax.inference.rauch_tung_stribel_smoother ¶

rauch_tung_stribel_smoother(transition_matrix_fn, t0, t1, mu0_s, cov0_s, mu0, cov0, mu0_, cov0_)

Discrete time Rauch-Tung-Striebel smoothing.

Parameters:

Name	Type	Description	Default
`transition_matrix_fn`	`Callable`	Transition matrix function	required
`t0`	`float`	Time of start	required
`t1`	`float`	Time of end	required
`mu0_s`	`Array`	Smoothed mean at t0	required
`cov0_s`	`Array`	Smoothed covariance at t0	required
`mu0`	`Array`	Unsmoothed mean at t0	required
`cov0`	`Array`	Unsmoothed covariance at t0	required
`mu0_`	`Array`	Prediction mean at t0	required
`cov0_`	`Array`	Prediction covariance at t0	required

Returns:

Type	Description
`Tuple[Array, Array]`	Tuple[Array, Array]: Updated mean and covariance.

Source code in probjax/inference/filtering/smoothing.py

def rauch_tung_stribel_smoother(
    transition_matrix_fn: Callable,
    t0: float,
    t1: float,
    mu0_s: Array,
    cov0_s: Array,
    mu0: Array,
    cov0: Array,
    mu0_: Array,
    cov0_: Array,
) -> Tuple[Array, Array]:
    """Discrete time Rauch-Tung-Striebel smoothing.

    Args:
        transition_matrix_fn (Callable): Transition matrix function
        t0 (float): Time of start
        t1 (float): Time of end
        mu0_s (Array): Smoothed mean at t0
        cov0_s (Array): Smoothed covariance at t0
        mu0 (Array): Unsmoothed mean at t0
        cov0 (Array): Unsmoothed covariance at t0
        mu0_ (Array): Prediction mean at t0
        cov0_ (Array): Prediction covariance at t0

    Returns:
        Tuple[Array, Array]: Updated mean and covariance.
    """
    Phi = transition_matrix_fn(t0, t1)
    G = jnp.dot(cov0, jnp.linalg.solve(cov0_, Phi).T)
    mu1 = mu0 + jnp.dot(G, mu0_s - jnp.dot(Phi, mu0_))
    cov1 = cov0 + jnp.dot(G, jnp.dot(cov0_s - cov0_, G.T))
    return mu1, cov1

probjax.inference.smooth ¶

smooth(ts, mus, covs, mus_, covs_, smooth)

Smooths the state given a Kalman filter output.

Parameters:

Name	Type	Description	Default
`ts`	`Array`	Time grid	required
`mus`	`Array`	Means	required
`covs`	`Array`	Covs	required
`mus_`	`Array`	Predicted means	required
`covs_`	`Array`	Predicted covs	required
`smooth`	`Callable`	Smoothing function	required

Returns:

Type	Description
`Tuple[Array, Array]`	Tuple[Array, Array]: description

Source code in probjax/inference/filtering/smoothing.py

def smooth(
    ts: Array, mus: Array, covs: Array, mus_: Array, covs_: Array, smooth: Callable
) -> Tuple[Array, Array]:
    """Smooths the state given a Kalman filter output.

    Args:
        ts (Array): Time grid
        mus (Array): Means
        covs (Array): Covs
        mus_ (Array): Predicted means
        covs_ (Array): Predicted covs
        smooth (Callable): Smoothing function

    Returns:
        Tuple[Array, Array]: _description_
    """

    idx_last = jnp.where((mus != mus_).all(-1))[-1][-1]
    mus_needed_ = jnp.flip(mus_[1 : idx_last + 1])
    covs_needed_ = jnp.flip(covs_[1 : idx_last + 1])
    mus_needed = jnp.flip(mus[:idx_last])
    covs_needed = jnp.flip(covs[:idx_last])
    ts_needed = jnp.flip(ts[:idx_last])

    def scan_fun(carry, data):
        (mu0_s, cov0_s, t1) = carry
        t0, mu0, cov0, mu0_, cov0_ = data
        mu1, cov1 = smooth(t0, t1, mu0_s, cov0_s, mu0, cov0, mu0_, cov0_)
        return (mu1, cov1, t0), (mu1, cov1)

    init_carry = (mus[idx_last], covs[idx_last], ts[idx_last])
    _, (mus_s, covs_s) = jax.lax.scan(
        scan_fun,
        init_carry,
        (ts_needed, mus_needed, covs_needed, mus_needed_, covs_needed_),
    )

    mus = jnp.concatenate([mus_s[::-1], mus[idx_last:]])
    covs = jnp.concatenate([covs_s[::-1], covs[idx_last:]])

    return mus, covs

Variational inference¶

probjax.inference.flow_vi ¶

Variational inference with a normalizing flow as the variational family.

Shaped after :mod:blackjax.vi.meanfield_vi -- same init/step/sample layout, same :class:~blackjax.base.VIAlgorithm return type, same stl_estimator option -- so it reads like the rest of the inference stack. The difference is the family: a mean-field Gaussian cannot represent a curved or correlated posterior, and a flow can.

The objective is the reparameterised reverse KL, mean(log q(x) - log p(x)) over samples x drawn from the flow.

Reverse KL is mode-seeking. A flow fitted this way tends to under-cover: on Neal's funnel it concentrates in the neck and reports a smaller variance than the truth. That is a property of the objective, not of this implementation, and it is the reason :func:probjax.inference.neutra exists -- running MCMC in the flow's latent space stays asymptotically exact however imperfect the flow is, so the flow only has to be helpful, not correct.

algorithm = flow_vi(logdensity_fn, flow, optax.adam(1e-3)) state = algorithm.init() def one(state, key): ... state, info = algorithm.step(key, state) ... return state, info.elbo state, objective = jax.lax.scan(one, state, jax.random.split(key, 2000)) draws = algorithm.sample(key, state, 1000)

FlowVIState ¶

Bases: NamedTuple

Variational parameters and the optimizer state that drives them.

Source code in probjax/inference/vi/flow_vi.py

class FlowVIState(NamedTuple):
    """Variational parameters and the optimizer state that drives them."""

    flow_params: nnx.State
    opt_state: OptState

FlowVIInfo ¶

Bases: NamedTuple

Per-step diagnostics.

Attributes:

Name	Type	Description
`elbo`	`float`	the value of the minimised objective, `mean(log q - log p)`. Despite the name -- kept for parity with `blackjax`'s `MFVIInfo` -- this is the negative ELBO shifted by the target's unknown log-normaliser, so it is the quantity that should go down during a run. Its absolute value is not interpretable; its trend is.

Source code in probjax/inference/vi/flow_vi.py

class FlowVIInfo(NamedTuple):
    """Per-step diagnostics.

    Attributes:
        elbo: the value of the minimised objective, ``mean(log q - log p)``.
            Despite the name -- kept for parity with ``blackjax``'s
            ``MFVIInfo`` -- this is the *negative* ELBO shifted by the target's
            unknown log-normaliser, so it is the quantity that should go **down**
            during a run. Its absolute value is not interpretable; its trend is.
    """

    elbo: float

init ¶

init(flow, optimizer)

Initialise from a flow, which supplies the variational family.

Unlike blackjax.vi.meanfield_vi.init there is no position argument: a mean-field family is defined by the shape of a position, whereas a flow already carries its own event size and initial parameters.

Source code in probjax/inference/vi/flow_vi.py

def init(flow, optimizer: GradientTransformation) -> FlowVIState:
    """Initialise from a flow, which supplies the variational family.

    Unlike ``blackjax.vi.meanfield_vi.init`` there is no ``position`` argument:
    a mean-field family is defined by the shape of a position, whereas a flow
    already carries its own event size and initial parameters.
    """
    _, flow_params, _ = _split(flow)
    return FlowVIState(flow_params, optimizer.init(flow_params))

step ¶

step(rng_key, state, logdensity_fn, optimizer, graphdef, rest, event_dim, num_samples=100, stl_estimator=True)

One reparameterised reverse-KL step.

graphdef/rest/event_dim come from splitting the flow once, which :func:as_top_level_api does for you; they are arguments rather than closure state so this mirrors blackjax, where step takes its configuration explicitly.

Source code in probjax/inference/vi/flow_vi.py

def step(
    rng_key,
    state: FlowVIState,
    logdensity_fn: Callable,
    optimizer: GradientTransformation,
    graphdef,
    rest,
    event_dim: int,
    num_samples: int = 100,
    stl_estimator: bool = True,
) -> tuple[FlowVIState, FlowVIInfo]:
    """One reparameterised reverse-KL step.

    ``graphdef``/``rest``/``event_dim`` come from splitting the flow once, which
    :func:`as_top_level_api` does for you; they are arguments rather than closure
    state so this mirrors blackjax, where ``step`` takes its configuration
    explicitly.
    """

    def objective(flow_params):
        samples = _draw(graphdef, flow_params, rest, rng_key, event_dim, num_samples)
        # Sticking the landing: with the entropy term's parameters detached the
        # score-function part of the gradient drops out, and it is exactly zero
        # in expectation -- so this removes variance without adding bias.
        entropy_params = (
            jax.lax.stop_gradient(flow_params) if stl_estimator else flow_params
        )
        log_q = _log_q(graphdef, entropy_params, rest, samples)
        log_p = jax.vmap(logdensity_fn)(samples)
        return jnp.mean(log_q - log_p)

    value, gradient = jax.value_and_grad(objective)(state.flow_params)
    updates, opt_state = optimizer.update(gradient, state.opt_state, state.flow_params)
    flow_params = optax.apply_updates(state.flow_params, updates)
    return FlowVIState(flow_params, opt_state), FlowVIInfo(value)

sample ¶

sample(rng_key, state, graphdef, rest, event_dim, num_samples=1)

Draw from the fitted approximation.

Source code in probjax/inference/vi/flow_vi.py

def sample(
    rng_key,
    state: FlowVIState,
    graphdef,
    rest,
    event_dim: int,
    num_samples: int = 1,
):
    """Draw from the fitted approximation."""
    return _draw(graphdef, state.flow_params, rest, rng_key, event_dim, num_samples)

as_top_level_api ¶

as_top_level_api(logdensity_fn, flow, optimizer, num_samples=100, stl_estimator=True)

Variational inference with a normalizing flow.

Parameters:

Name	Type	Description	Default
`logdensity_fn`	`Callable`	the unnormalized target log-density, taking one position.	required
`flow`	`NormalizingFlow`	a flow from :mod:`probjax.nn` -- `maf`, `nsf`, `naf` and the rest all work. Its parameters are the variational parameters; the flow itself is not mutated.	required
`optimizer`	`GradientTransformation`	an optax `GradientTransformation`.	required
`num_samples`	`int`	draws used to estimate the objective at each step.	`100`
`stl_estimator`	`bool`	use the sticking-the-landing gradient estimator, which lowers gradient variance at no cost in bias.	`True`

Returns:

Type	Description
`VIAlgorithm`	A `blackjax` :class:`~blackjax.base.VIAlgorithm`: `init`, `step`
`VIAlgorithm`	and `sample`. As with blackjax, the loop over `step` is the caller's
`VIAlgorithm`	-- see the module docstring for the `lax.scan` form.

Source code in probjax/inference/vi/flow_vi.py

def as_top_level_api(
    logdensity_fn: Callable,
    flow,
    optimizer: GradientTransformation,
    num_samples: int = 100,
    stl_estimator: bool = True,
) -> VIAlgorithm:
    """Variational inference with a normalizing flow.

    Args:
        logdensity_fn: the unnormalized target log-density, taking one position.
        flow (NormalizingFlow): a flow from :mod:`probjax.nn` -- ``maf``,
            ``nsf``, ``naf`` and the rest all work. Its parameters are the
            variational parameters; the flow itself is not mutated.
        optimizer: an optax ``GradientTransformation``.
        num_samples: draws used to estimate the objective at each step.
        stl_estimator: use the sticking-the-landing gradient estimator, which
            lowers gradient variance at no cost in bias.

    Returns:
        A ``blackjax`` :class:`~blackjax.base.VIAlgorithm`: ``init``, ``step``
        and ``sample``. As with blackjax, the loop over ``step`` is the caller's
        -- see the module docstring for the ``lax.scan`` form.
    """
    graphdef, _, rest = _split(flow)
    event_dim = flow.input_dim

    def init_fn():
        return init(flow, optimizer)

    def step_fn(rng_key, state: FlowVIState) -> tuple[FlowVIState, FlowVIInfo]:
        return step(
            rng_key,
            state,
            logdensity_fn,
            optimizer,
            graphdef,
            rest,
            event_dim,
            num_samples,
            stl_estimator,
        )

    def sample_fn(rng_key, state: FlowVIState, num_samples: int = 1):
        return sample(rng_key, state, graphdef, rest, event_dim, num_samples)

    return VIAlgorithm(init_fn, step_fn, sample_fn)

rebuild ¶

rebuild(flow, state)

Return a flow carrying the fitted variational parameters.

The bridge to :func:probjax.inference.neutra, which wants a flow rather than a parameter tree.

Source code in probjax/inference/vi/flow_vi.py

def rebuild(flow, state: FlowVIState):
    """Return a flow carrying the fitted variational parameters.

    The bridge to :func:`probjax.inference.neutra`, which wants a flow rather
    than a parameter tree.
    """
    graphdef, _, rest = _split(flow)
    return nnx.merge(graphdef, state.flow_params, rest)

probjax.inference.neutra ¶

neutra(logdensity_fn, flow)

Reparameterise logdensity_fn through flow.

Parameters:

Name	Type	Description	Default
`logdensity_fn`	`Callable`	the unnormalized target log-density, taking one position.	required
`flow`	`NormalizingFlow`	a normalizing flow, typically fitted with :func:`probjax.inference.flow_vi`, but any flow works -- one trained on posterior samples from a previous run is equally valid.	required

Returns:

Name	Type	Description
`A`	`NeuTraTransform`	class:`NeuTraTransform`.

Source code in probjax/inference/vi/neutra.py

def neutra(logdensity_fn: Callable, flow) -> NeuTraTransform:
    """Reparameterise ``logdensity_fn`` through ``flow``.

    Args:
        logdensity_fn: the unnormalized target log-density, taking one position.
        flow (NormalizingFlow): a normalizing flow, typically fitted with
            :func:`probjax.inference.flow_vi`, but any flow works -- one trained
            on posterior samples from a previous run is equally valid.

    Returns:
        A :class:`NeuTraTransform`.
    """
    log_q = flow._logpdf

    def transformed_logdensity(position):
        target_position = flow.transform(position)
        # log|det J_T(z)| is not exposed by the flow, but it does not need to be:
        #     log q(T(z)) = log N(z) - log|det J_T(z)|
        # so the Jacobian term is the difference between the base density at z
        # and the flow's own density at T(z). Both are one call, and it is exact
        # -- no Jacobian is ever formed.
        log_base = jnp.sum(jax.scipy.stats.norm.logpdf(position))
        log_jacobian = log_base - jnp.squeeze(log_q(target_position))
        return logdensity_fn(target_position) + log_jacobian

    def forward(positions):
        positions = jnp.asarray(positions)
        if positions.ndim == 1:
            return flow.transform(positions)
        flat = positions.reshape(-1, positions.shape[-1])
        mapped = jax.vmap(flow.transform)(flat)
        return mapped.reshape(*positions.shape[:-1], mapped.shape[-1])

    return NeuTraTransform(transformed_logdensity, forward)

probjax.inference.FlowVIState ¶

Bases: NamedTuple

Variational parameters and the optimizer state that drives them.

Source code in probjax/inference/vi/flow_vi.py

class FlowVIState(NamedTuple):
    """Variational parameters and the optimizer state that drives them."""

    flow_params: nnx.State
    opt_state: OptState

probjax.inference.FlowVIInfo ¶

Bases: NamedTuple

Per-step diagnostics.

Attributes:

Name	Type	Description
`elbo`	`float`	the value of the minimised objective, `mean(log q - log p)`. Despite the name -- kept for parity with `blackjax`'s `MFVIInfo` -- this is the negative ELBO shifted by the target's unknown log-normaliser, so it is the quantity that should go down during a run. Its absolute value is not interpretable; its trend is.

Source code in probjax/inference/vi/flow_vi.py

class FlowVIInfo(NamedTuple):
    """Per-step diagnostics.

    Attributes:
        elbo: the value of the minimised objective, ``mean(log q - log p)``.
            Despite the name -- kept for parity with ``blackjax``'s
            ``MFVIInfo`` -- this is the *negative* ELBO shifted by the target's
            unknown log-normaliser, so it is the quantity that should go **down**
            during a run. Its absolute value is not interpretable; its trend is.
    """

    elbo: float

probjax.inference.NeuTraTransform ¶

Bases: NamedTuple

A target rewritten in a flow's latent coordinates.

Attributes:

Name	Type	Description
`logdensity`	`Callable`	the pulled-back log-density, to hand to any kernel.
`forward`	`Callable`	maps latent draws back to the target's space. Accepts a single position or a leading batch of them.

Source code in probjax/inference/vi/neutra.py

class NeuTraTransform(NamedTuple):
    """A target rewritten in a flow's latent coordinates.

    Attributes:
        logdensity: the pulled-back log-density, to hand to any kernel.
        forward: maps latent draws back to the target's space. Accepts a single
            position or a leading batch of them.
    """

    logdensity: Callable
    forward: Callable

States and results¶

The types returned by the runners and kernels. Kernel State and Params are blackjax's own types, re-exported for convenience and documented there.

probjax.inference.MCMCResult ¶

Bases: NamedTuple, Generic[S, I]

Source code in probjax/inference/base.py

class MCMCResult(NamedTuple, Generic[S, I]):
    state: S
    samples: Optional[Any] = None
    info: Optional[I] = None

probjax.inference.SMCResult ¶

Bases: NamedTuple, Generic[S, P, I]

Source code in probjax/inference/base.py

class SMCResult(NamedTuple, Generic[S, P, I]):
    state: S
    params: P
    info: Optional[I] = None
    log_evidence: Optional[Any] = None
    final_info: Optional[I] = None
    num_steps: Optional[Any] = None
    completed: Optional[Any] = None
    key: Optional[Any] = None

probjax.inference.FilteringResult ¶

Bases: NamedTuple, Generic[S, I]

Source code in probjax/inference/base.py

class FilteringResult(NamedTuple, Generic[S, I]):
    initial_state: S
    states: Any
    info: Optional[I] = None

probjax.inference.AdaptationResult ¶

Bases: NamedTuple, Generic[S, P, I]

Source code in probjax/inference/base.py

class AdaptationResult(NamedTuple, Generic[S, P, I]):
    state: S
    params: P
    trace: Optional[Any] = None
    final_info: Optional[I] = None

probjax.inference.WarmupResult ¶

Bases: NamedTuple, Generic[S, P, I]

Source code in probjax/inference/base.py

class WarmupResult(NamedTuple, Generic[S, P, I]):
    state: S
    params: P
    info: Optional[I] = None

probjax.inference.MarkovKernel `module-attribute` ¶

MarkovKernel = Kernel

probjax.inference.Kernel ¶

Bases: NamedTuple

A pure transition and the functions needed to initialize it.

default_adaptation is an optional factory returning the kernel's recommended adaptation policy (an Adaptor or a Warmup). Runners use it when no explicit policy is passed; an explicitly passed policy always wins.

Source code in probjax/inference/base.py

class Kernel(NamedTuple):
    """A pure transition and the functions needed to initialize it.

    ``default_adaptation`` is an optional factory returning the kernel's
    recommended adaptation policy (an ``Adaptor`` or a ``Warmup``).
    Runners use it when no explicit policy is passed; an explicitly
    passed policy always wins.
    """

    init: Callable
    step: Callable
    init_params: Callable
    default_adaptation: Optional[Callable[..., Any]] = None

    def __call__(self, key, state, params=None, *args, **kwargs):
        return self.step(key, state, params, *args, **kwargs)

probjax.inference.Warmup ¶

Bases: NamedTuple

A finite initialization policy for a kernel and its parameters.

Source code in probjax/inference/base.py

class Warmup(NamedTuple):
    """A finite initialization policy for a kernel and its parameters."""

    run: Callable

probjax.inference.Adaptor ¶

Bases: NamedTuple

A local parameter-adaptation state machine.

init receives (state, params). update consumes one completed transition and can therefore be used during warmup or regular sampling.

Source code in probjax/inference/base.py

class Adaptor(NamedTuple):
    """A local parameter-adaptation state machine.

    ``init`` receives ``(state, params)``. ``update`` consumes one completed
    transition and can therefore be used during warmup or regular sampling.
    """

    init: Callable
    update: Callable
    finalize: Callable

probjax.inference.FilterState ¶

Bases: NamedTuple

This is a NamedTuple that represents the state of a filter.

It contains all the information required to run the filter.

Source code in probjax/inference/filtering/base.py

class FilterState(NamedTuple):
    """This is a NamedTuple that represents the state of a filter.

    It contains all the information **required** to run the filter.
    """

    pass

probjax.inference.FilterInfo ¶

Bases: NamedTuple

This is a NamedTuple that represents the information returned by a filter.

It contains all useful information that can be extracted from the filter.

Source code in probjax/inference/filtering/base.py

class FilterInfo(NamedTuple):
    """This is a NamedTuple that represents the information returned by a filter.

    It contains all useful information that can be extracted from the filter.

    """

    pass

probjax.inference.FilterKernel ¶

Bases: NamedTuple

This is a NamedTuple that represents a filter kernel.

Source code in probjax/inference/filtering/base.py

class FilterKernel(NamedTuple):
    """This is a NamedTuple that represents a filter kernel."""

    init: Callable
    step: Callable
    default_unpack: Callable = _default_unpack

    def __call__(
        self,
        state: FilterState,
        t: Optional[ArrayLike] = None,
        observed: Optional[ArrayLike] = None,
        rng_key: Optional[ArrayLike] = None,
    ) -> Tuple[FilterState, FilterInfo]:
        return self.step(state, t, observed, rng_key)

Inference¶

Runners¶

probjax.inference.MCMC ¶

run_kernel staticmethod ¶

run ¶

sample ¶

adapt ¶

warmup ¶

probjax.inference.SMC ¶

run_adaptive_kernel staticmethod ¶

MCMC kernels¶

probjax.inference.hmc ¶

probjax.inference.nuts ¶

probjax.inference.dynamic_hmc ¶

probjax.inference.mala ¶

probjax.inference.mclmc ¶

probjax.inference.adjusted_mclmc ¶

probjax.inference.mh ¶

probjax.inference.gauss_rwmh ¶

probjax.inference.imh ¶

probjax.inference.neural_imh ¶

probjax.inference.slice ¶

probjax.inference.latent_slice ¶

probjax.inference.elliptical_slice ¶

probjax.inference.arms ¶

probjax.inference.a2rms ¶

probjax.inference.ars ¶

probjax.inference.gaussian_imh ¶

probjax.inference.adjusted_mclmc_dynamic ¶

probjax.inference.RejectionSampler ¶

probjax.inference.pseudo_marginal ¶

Stochastic-gradient MCMC¶

probjax.inference.sgld ¶

probjax.inference.sghmc ¶

probjax.inference.sgnht ¶

Warmup and adaptation¶

probjax.inference.window_warmup ¶

probjax.inference.pathfinder_warmup ¶

probjax.inference.mclmc_warmup ¶

probjax.inference.neural_imh_warmup ¶

probjax.inference.step_size_adaptor ¶

probjax.inference.mass_matrix_adaptor ¶

probjax.inference.covariance_adaptor ¶

probjax.inference.mean_adaptor ¶

probjax.inference.compose_adaptors ¶

probjax.inference.acceptance_rate_adaptor ¶

probjax.inference.slice_step_size_adaptor ¶

probjax.inference.particle_adaptor ¶

probjax.inference.adapt module-attribute ¶

probjax.inference.adapt_step ¶

probjax.inference.as_warmup ¶

Convergence diagnostics¶

probjax.inference.effective_sample_size ¶

Parameters¶

Returns¶

Notes¶

probjax.inference.ess_bulk ¶

Parameters¶

Returns¶

Notes¶

References¶

probjax.inference.ess_tail ¶

Parameters¶

Returns¶

Notes¶

References¶

probjax.inference.potential_scale_reduction ¶

Parameters¶

Returns¶

Notes¶

probjax.inference.rhat ¶

Parameters¶

Returns¶

Notes¶

References¶

probjax.inference.pareto_khat ¶

Parameters¶

Returns¶

Notes¶

probjax.inference.kernelized_stein_discrepancy ¶

run_kernel `staticmethod` ¶

run_adaptive_kernel `staticmethod` ¶

probjax.inference.adapt `module-attribute` ¶

probjax.inference.GeometricPath `dataclass` ¶

probjax.inference.PartialPosteriorsPath `dataclass` ¶