Inference¶

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

Runners¶

probjax.inference.MCMC ¶

Bases: WithProgressBarAPI

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):
    """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)

        def one_step(state, xs):
            if args is None:
                step_key, step_args = xs, ()
            else:
                step_key, step_args = xs
            state, info = kernel(step_key, state, params, *step_args)
            return state, (
                _select(state, collect_state),
                _select(info, 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)

    @staticmethod
    @partial(
        jax.jit,
        static_argnames=(
            "kernel",
            "num_samples",
            "thin",
            "collect_info",
        ),
    )
    def sample_kernel(
        key: RngKey,
        kernel: MarkovKernel,
        state: State,
        num_samples: int,
        params: Params,
        args: Optional[Tuple] = None,
        *,
        thin: int = 1,
        collect_info: Tuple[str, ...] = (),
    ) -> MCMCResult:
        """Collect positions from a compiled kernel scan."""
        keys = jax.random.split(key, (num_samples, thin))
        step_args = None
        if args is not None:
            step_args = jax.tree.map(
                lambda value: value.reshape((num_samples, thin) + value.shape[1:]),
                args,
            )
        xs = keys if step_args is None else (keys, step_args)

        def collect_one(state, xs):
            if step_args is None:
                sample_keys, sample_args = xs, None
            else:
                sample_keys, sample_args = xs
            inner_xs = (
                sample_keys if sample_args is None else (sample_keys, sample_args)
            )

            def transition(state, inner_xs):
                if sample_args is None:
                    step_key, current_args = inner_xs, ()
                else:
                    step_key, current_args = inner_xs
                state, info = kernel(step_key, state, params, *current_args)
                return state, (info, _select(info, collect_info))

            state, (step_info, selected_info) = jax.lax.scan(
                transition, state, inner_xs
            )
            diagnostics = jax.tree.map(lambda value: value[-1], selected_info)
            return state, (state.position, diagnostics)

        state, (samples, info) = jax.lax.scan(collect_one, state, xs)
        return MCMCResult(state, samples, info if collect_info else None)

    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 _verbose_scan(self, f, init, xs, length, stats_fn):
        """Run a scan with a rate-limited progress bar.

        ``print_scan`` requires a tuple carry, so we wrap the kernel state.
        """

        def wrapped(carry, x):
            state, y = f(carry[0], x)
            return (state,), y

        update_stats, print_fn, init_stats, print_rate = self._make_verbose_fns(
            length, stats_fn=lambda carry, y: stats_fn(carry[0], y)
        )
        (state,), y = print_scan(
            wrapped,
            (init,),
            init_stats,
            xs=xs,
            length=length,
            update_stats=update_stats,
            print_rate=print_rate,
            print_fn=print_fn,
        )
        return state, y

    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)

        def one_step(state, xs):
            if args is None:
                step_key, step_args = xs, ()
            else:
                step_key, step_args = xs
            state, info = self.kernel(step_key, state, params, *step_args)
            return state, (
                self._extract_stats(state, info),
                _select(state, self.collect_state),
                _select(info, self.collect_info),
            )

        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 = None
        if args is not None:
            step_args = jax.tree.map(
                lambda value: value.reshape((num_samples, thin) + value.shape[1:]),
                args,
            )
        xs = keys if step_args is None else (keys, step_args)

        def collect_one(state, xs):
            if step_args is None:
                sample_keys, sample_args = xs, None
            else:
                sample_keys, sample_args = xs
            inner_xs = (
                sample_keys if sample_args is None else (sample_keys, sample_args)
            )

            def transition(state, inner_xs):
                if sample_args is None:
                    step_key, current_args = inner_xs, ()
                else:
                    step_key, current_args = inner_xs
                state, info = self.kernel(step_key, state, params, *current_args)
                return state, (info, _select(info, self.collect_info))

            state, (step_info, selected_info) = jax.lax.scan(
                transition, state, inner_xs
            )
            last_info = jax.tree.map(lambda value: value[-1], step_info)
            diagnostics = jax.tree.map(lambda value: value[-1], selected_info)
            return state, (
                self._extract_stats(state, last_info),
                state.position,
                diagnostics,
            )

        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 warmup_kernel(
        key: RngKey,
        kernel: MarkovKernel,
        warmup: Warmup,
        state: State,
        params: Params,
        num_steps: int,
    ) -> WarmupResult:
        return warmup.run(key, kernel, state, params, num_steps)

    def warmup(
        self,
        key: RngKey,
        warmup: Warmup,
        state: State,
        params: Params,
        num_steps: int,
    ) -> WarmupResult:
        """Run a specialized state-producing warmup procedure."""
        return self.warmup_kernel(key, self.kernel, warmup, 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)

    def one_step(state, xs):
        if args is None:
            step_key, step_args = xs, ()
        else:
            step_key, step_args = xs
        state, info = kernel(step_key, state, params, *step_args)
        return state, (
            _select(state, collect_state),
            _select(info, 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)

sample_kernel `staticmethod` ¶

sample_kernel(key, kernel, state, num_samples, params, args=None, *, thin=1, collect_info=())

Collect positions from a compiled kernel scan.

Source code in probjax/inference/mcmc_runner.py

@staticmethod
@partial(
    jax.jit,
    static_argnames=(
        "kernel",
        "num_samples",
        "thin",
        "collect_info",
    ),
)
def sample_kernel(
    key: RngKey,
    kernel: MarkovKernel,
    state: State,
    num_samples: int,
    params: Params,
    args: Optional[Tuple] = None,
    *,
    thin: int = 1,
    collect_info: Tuple[str, ...] = (),
) -> MCMCResult:
    """Collect positions from a compiled kernel scan."""
    keys = jax.random.split(key, (num_samples, thin))
    step_args = None
    if args is not None:
        step_args = jax.tree.map(
            lambda value: value.reshape((num_samples, thin) + value.shape[1:]),
            args,
        )
    xs = keys if step_args is None else (keys, step_args)

    def collect_one(state, xs):
        if step_args is None:
            sample_keys, sample_args = xs, None
        else:
            sample_keys, sample_args = xs
        inner_xs = (
            sample_keys if sample_args is None else (sample_keys, sample_args)
        )

        def transition(state, inner_xs):
            if sample_args is None:
                step_key, current_args = inner_xs, ()
            else:
                step_key, current_args = inner_xs
            state, info = kernel(step_key, state, params, *current_args)
            return state, (info, _select(info, collect_info))

        state, (step_info, selected_info) = jax.lax.scan(
            transition, state, inner_xs
        )
        diagnostics = jax.tree.map(lambda value: value[-1], selected_info)
        return state, (state.position, diagnostics)

    state, (samples, info) = jax.lax.scan(collect_one, state, xs)
    return MCMCResult(state, samples, info if collect_info else None)

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)

    def one_step(state, xs):
        if args is None:
            step_key, step_args = xs, ()
        else:
            step_key, step_args = xs
        state, info = self.kernel(step_key, state, params, *step_args)
        return state, (
            self._extract_stats(state, info),
            _select(state, self.collect_state),
            _select(info, self.collect_info),
        )

    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 = None
    if args is not None:
        step_args = jax.tree.map(
            lambda value: value.reshape((num_samples, thin) + value.shape[1:]),
            args,
        )
    xs = keys if step_args is None else (keys, step_args)

    def collect_one(state, xs):
        if step_args is None:
            sample_keys, sample_args = xs, None
        else:
            sample_keys, sample_args = xs
        inner_xs = (
            sample_keys if sample_args is None else (sample_keys, sample_args)
        )

        def transition(state, inner_xs):
            if sample_args is None:
                step_key, current_args = inner_xs, ()
            else:
                step_key, current_args = inner_xs
            state, info = self.kernel(step_key, state, params, *current_args)
            return state, (info, _select(info, self.collect_info))

        state, (step_info, selected_info) = jax.lax.scan(
            transition, state, inner_xs
        )
        last_info = jax.tree.map(lambda value: value[-1], step_info)
        diagnostics = jax.tree.map(lambda value: value[-1], selected_info)
        return state, (
            self._extract_stats(state, last_info),
            state.position,
            diagnostics,
        )

    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, warmup, state, params, num_steps)

Run a specialized state-producing warmup procedure.

Source code in probjax/inference/mcmc_runner.py

def warmup(
    self,
    key: RngKey,
    warmup: Warmup,
    state: State,
    params: Params,
    num_steps: int,
) -> WarmupResult:
    """Run a specialized state-producing warmup procedure."""
    return self.warmup_kernel(key, self.kernel, warmup, state, params, num_steps)

probjax.inference.SMC ¶

Bases: WithProgressBarAPI

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):
    """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.
        """
        if not isinstance(max_steps, int) or max_steps < 1:
            raise ValueError("max_steps must be a positive integer.")
        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 _verbose_scan(self, f, init, xs, length, stats_fn):
        def wrapped(carry, x):
            state, y = f(carry[0], x)
            return (state,), y

        update_stats, print_fn, init_stats, print_rate = self._make_verbose_fns(
            length, stats_fn=lambda carry, y: stats_fn(carry[0], y)
        )
        (state,), y = print_scan(
            wrapped,
            (init,),
            init_stats,
            xs,
            length,
            update_stats=update_stats,
            print_fn=print_fn,
            print_rate=print_rate,
        )
        return state, y

    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.
    """
    if not isinstance(max_steps, int) or max_steps < 1:
        raise ValueError("max_steps must be a positive integer.")
    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 ¶

init_params ¶

init_params(state, step_size=0.5, inverse_mass_matrix=None)

Initialize the parameters for the HMC kernel.

Parameters:

Name	Type	Description	Default
`state`	`PyTree`	Position of the chain.	required
`step_size`	`float`	Default step size for the HMC kernel. Defaults to 0.5.	`0.5`
`inverse_mass_matrix`	`Optional[Array]`	Inverse mass matrix. Defaults to None i.e. identity matrix (as diagonal!).	`None`

Raises:

Type	Description
`ValueError`	If the dimension of the inverse mass matrix does not match the dimension of the position.

Returns:

Name	Type	Description
`HMCParams`	`HMCParams`	Parameters for the HMC kernel.

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

def init_params(
    state: PyTree,
    step_size: float = 0.5,
    inverse_mass_matrix: Optional[Array] = None,
) -> HMCParams:
    """Initialize the parameters for the HMC kernel.

    Args:
        state (PyTree): Position of the chain.
        step_size (float): Default step size for the HMC kernel. Defaults to 0.5.
        inverse_mass_matrix (Optional[Array], optional): Inverse mass matrix.
            Defaults to None i.e. identity matrix (as diagonal!).

    Raises:
        ValueError: If the dimension of the inverse mass matrix does not match
            the dimension of the position.

    Returns:
        HMCParams: Parameters for the HMC kernel.
    """
    inverse_mass_matrix = _init_hmc_like_params(state, inverse_mass_matrix)
    step_size = _scale_step_size_by_grad(state, step_size, inverse_mass_matrix)
    return HMCParams(
        step_size=step_size,
        inverse_mass_matrix=inverse_mass_matrix,
    )

build_step ¶

build_step(logdensity_fn, num_integration_steps=10, integrator=velocity_verlet, divergence_threshold=1000.0)

Build the HMC kernel.

Parameters:

Name	Type	Description	Default
`logdensity_fn`	`Callable`	The log density function.	required
`num_integration_steps`	`int`	Number of integration steps. Defaults to 10.	`10`
`integrator`	`Callable`	The integrator to use. Defaults to blackjax.integrators.velocity_verlet.	`velocity_verlet`
`divergence_threshold`	`float`	The threshold for the divergence check. Defaults to 1000.0.	`1000.0`

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

def build_step(
    logdensity_fn: Callable,
    num_integration_steps: int = 10,
    integrator: Callable = blackjax.mcmc.integrators.velocity_verlet,
    divergence_threshold: float = 1000.0,
):
    """Build the HMC kernel.

    Args:
        logdensity_fn (Callable): The log density function.
        num_integration_steps (int, optional): Number of integration steps.
            Defaults to 10.
        integrator (Callable, optional): The integrator to use. Defaults to
            blackjax.integrators.velocity_verlet.
        divergence_threshold (float, optional): The threshold for the divergence
            check. Defaults to 1000.0.

    """
    kernel_builder = lambda: blackjax.hmc.build_kernel(integrator, divergence_threshold)
    return make_step_from_kernel(
        logdensity_fn,
        kernel_builder,
        call_defaults={"num_integration_steps": num_integration_steps},
    )

probjax.inference.nuts `module-attribute` ¶

nuts = make_kernel_api(name='nuts', init_fn=blackjax.hmc.init, init_params_fn=init_params, build_step_fn=build_kernel_nuts)

probjax.inference.dynamic_hmc ¶

probjax.inference.mala ¶

probjax.inference.mclmc ¶

probjax.inference.adjusted_mclmc `module-attribute` ¶

adjusted_mclmc = make_kernel_api(name='adjusted_mclmc', init_fn=blackjax.adjusted_mclmc.init, init_params_fn=init_params, build_step_fn=build_step)

probjax.inference.mh ¶

probjax.inference.gauss_rwmh `module-attribute` ¶

gauss_rwmh = make_kernel_api(name='gauss_rwmh', init_fn=blackjax.rmh.init, init_params_fn=init_params_gaussian_rw, build_step_fn=partial(build_mh_step, transition_proposal_fn=gaussian_transition_proposal))

probjax.inference.imh ¶

NeuralIMHParams ¶

Bases: NamedTuple

Parameters for an independent neural proposal.

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

class NeuralIMHParams(NamedTuple):
    """Parameters for an independent neural proposal."""

    proposal_state: PyTree
    iteration: Array

NeuralIMHState ¶

Bases: NamedTuple

IMH state with a cached proposal density.

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

class NeuralIMHState(NamedTuple):
    """IMH state with a cached proposal density."""

    position: PyTree
    logdensity: Array
    proposal_logdensity: Array
    proposal_iteration: Array

NeuralIMHWarmupInfo ¶

Bases: NamedTuple

Per-round diagnostics from neural proposal adaptation.

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

class NeuralIMHWarmupInfo(NamedTuple):
    """Per-round diagnostics from neural proposal adaptation."""

    losses: Array
    acceptance_rate: Array

GaussianIMHParams ¶

Bases: NamedTuple

Parameters for the Gaussian IMH kernel.

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

class GaussianIMHParams(NamedTuple):
    """Parameters for the Gaussian IMH kernel."""

    mean: Array
    cov: Array
    unflatten: Callable

wrap_logpdf ¶

wrap_logpdf(logpdf)

Wrap the logpdf function to work with the IMH kernel.

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

def wrap_logpdf(logpdf: Callable) -> Callable:
    """Wrap the logpdf function to work with the IMH kernel."""

    def wrapped_logpdf(x: Any, y: Any, *args, **kwargs) -> Array:
        return logpdf(y, *args, **kwargs)

    return wrapped_logpdf

build_imh_step ¶

build_imh_step(logdensity_fn, proposal_fn, proposal_logpdf)

A function to build the Independent Metropolis-Hastings kernel.

Parameters:

Name	Type	Description	Default
`logdensity_fn`	`Callable`	The log density function.	required
`proposal_fn`	`Callable`	Proposal function.	required
`proposal_logpdf`	`Callable`	Proposal log pdf.	required

Returns:

Name	Type	Description
`Callable`	`Callable`	Step function for the IMH kernel.

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

def build_imh_step(
    logdensity_fn: Callable,
    proposal_fn: Callable,
    proposal_logpdf: Callable,
) -> Callable:
    """A function to build the Independent Metropolis-Hastings kernel.

    Args:
        logdensity_fn (Callable): The log density function.
        proposal_fn (Callable): Proposal function.
        proposal_logpdf (Callable): Proposal log pdf.

    Returns:
        Callable: Step function for the IMH kernel.
    """
    kernel = blackjax.irmh.build_kernel()

    def step(
        key: RngKey, state: RWState, params: Optional[IMHParams] = None
    ) -> Tuple[RWState, RWInfo]:
        _proposal_fn = partial(proposal_fn, params=params)
        _proposal_logpdf = partial(wrap_logpdf(proposal_logpdf), params=params)
        return kernel(
            key,
            state,
            logdensity_fn,
            proposal_distribution=_proposal_fn,
            proposal_logdensity_fn=_proposal_logpdf,
        )

    return step

init_imh_params ¶

init_imh_params(state, rng_key=None)

Generally, there are no parameters to initialize for the IMH kernel.

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

def init_imh_params(state: PyTree, rng_key=None) -> IMHParams:
    """Generally, there are no parameters to initialize for the IMH kernel."""
    return IMHParams()

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)

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:
        import optax

        fit_optimizer = optax.adam(learning_rate)

    def run(key, kernel, state, params, num_steps):
        from flax import nnx

        from probjax.inference.mcmc_runner import MCMC

        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 = MCMC.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)

init_gaussian_imh_params ¶

init_gaussian_imh_params(state, mean=None, cov=None)

Initialize the parameters for the Gaussian IMH kernel.

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

def init_gaussian_imh_params(
    state: PyTree, mean: Optional[Array] = None, cov: Optional[Array] = None
) -> GaussianIMHParams:
    """Initialize the parameters for the Gaussian IMH kernel."""
    position = state.position if hasattr(state, "position") else state
    flat_position, unflatten = jax.flatten_util.ravel_pytree(position)
    if mean is None:
        mean = flat_position
    if cov is None:
        cov = jnp.ones_like(flat_position)
    return GaussianIMHParams(mean=mean, cov=cov, unflatten=unflatten)

proposal_gaussian ¶

proposal_gaussian(key, *, params)

Generate a new position from a Gaussian proposal.

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

def proposal_gaussian(key: RngKey, *, params: GaussianIMHParams):
    """Generate a new position from a Gaussian proposal."""
    mean = params.mean
    cov = params.cov
    eps = jax.random.normal(key, mean.shape)
    if cov.ndim == 1:
        new_position = mean + jnp.sqrt(cov) * eps
    else:
        new_position = mean + jnp.dot(jnp.linalg.cholesky(cov), eps)
    return params.unflatten(new_position)

proposal_gaussian_logpdf ¶

proposal_gaussian_logpdf(state, *, params)

Log pdf of the Gaussian proposal.

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

def proposal_gaussian_logpdf(state, *, params: GaussianIMHParams):
    """Log pdf of the Gaussian proposal."""
    x = state.position
    flat_position, _ = jax.flatten_util.ravel_pytree(x)
    mean = params.mean
    cov = params.cov
    if cov.ndim == 1:
        return jax.scipy.stats.norm.logpdf(flat_position, mean, cov).sum()
    else:
        return jax.scipy.stats.multivariate_normal.logpdf(flat_position, mean, cov)

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)

probjax.inference.slice ¶

probjax.inference.latent_slice ¶

probjax.inference.elliptical_slice ¶

probjax.inference.arms ¶

probjax.inference.a2rms `module-attribute` ¶

a2rms = make_kernel_api(name='a2rms', init_fn=init, init_params_fn=init_params, build_step_fn=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 `module-attribute` ¶

gaussian_imh = make_kernel_api(name='gaussian_imh', init_fn=blackjax.irmh.init, init_params_fn=init_gaussian_imh_params, build_step_fn=partial(build_imh_step, proposal_fn=proposal_gaussian, proposal_logpdf=proposal_gaussian_logpdf))

probjax.inference.adjusted_mclmc_dynamic `module-attribute` ¶

adjusted_mclmc_dynamic = make_kernel_api(name='adjusted_mclmc_dynamic', init_fn=init_dynamic, init_params_fn=init_dynamic_params, build_step_fn=build_dynamic_step)

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)

    return MarkovKernel(init, step, init_params)

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 = blackjax.sgld.build_kernel()

    def init(key, position=None, rng_key=None):
        if position is None:
            position = key
        return SGMCMCState(position=blackjax.sgld.init(position))

    def step(key: RngKey, state: SGMCMCState, params, *args):
        new_position = kernel(
            key,
            state.position,
            grad_estimator,
            minibatch=args[0] if args else None,
            step_size=params.step_size,
            temperature=params.temperature,
        )
        return SGMCMCState(position=new_position), SGMCMCInfo()

    return MarkovKernel(
        init,
        step,
        lambda state, step_size=1e-3, temperature=1.0: SGLDParams(
            step_size=step_size, temperature=temperature
        ),
    )

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.
    """
    kernel = blackjax.sghmc.build_kernel(alpha=alpha, beta=beta)

    def init(key, position=None, rng_key=None):
        if position is None:
            position = key
        return SGMCMCState(position=blackjax.sghmc.init(position))

    def step(key: RngKey, state: SGMCMCState, params, *args):
        new_position = kernel(
            key,
            state.position,
            grad_estimator,
            minibatch=args[0] if args else None,
            step_size=params.step_size,
            temperature=params.temperature,
            num_integration_steps=num_integration_steps,
        )
        return SGMCMCState(position=new_position), SGMCMCInfo()

    return MarkovKernel(
        init,
        step,
        lambda state, step_size=1e-3, temperature=1.0: SGHMCParams(
            step_size=step_size, temperature=temperature
        ),
    )

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 step(key: RngKey, state, params, *args):
        new_state = kernel(
            key,
            state,
            grad_estimator,
            minibatch=args[0] if args else None,
            step_size=params.step_size,
            temperature=params.temperature,
        )
        return new_state, SGMCMCInfo()

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

    return MarkovKernel(
        init,
        step,
        lambda state, step_size=1e-3, temperature=1.0: SGNHTParams(
            step_size=step_size, temperature=temperature
        ),
    )

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."""
    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:
        import optax

        fit_optimizer = optax.adam(learning_rate)

    def run(key, kernel, state, params, num_steps):
        from flax import nnx

        from probjax.inference.mcmc_runner import MCMC

        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 = MCMC.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.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."""

    def init(_state, _params):
        return ()

    def update(_state, info, adaptor_state, params):
        params = tune_from_kernel_info(
            _params_to_dict(params),
            info_fn(info),
            target_acceptance_rate=target_acceptance_rate,
            scale_key=scale_key,
        )
        return adaptor_state, _ensure_param_batch(params, shared=True), None

    def finalize(_adaptor_state, params):
        return params, None

    return Adaptor(init, update, finalize)

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."""

    def init(_state, _params):
        return ()

    def update(state, _info, adaptor_state, params):
        params = tune_from_particles(
            _params_to_dict(params), state.particles, updates=updates
        )
        return adaptor_state, _ensure_param_batch(params, shared=True), None

    def finalize(_adaptor_state, params):
        return params, None

    return Adaptor(init, update, finalize)

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."""
    from probjax.inference.base import Warmup

    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)

Sequential Monte Carlo¶

probjax.inference.smc ¶

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."""
    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.")

persistent_smc_kernel ¶

persistent_smc_kernel(**kwargs)

Build a persistent SMC kernel (geometric path only).

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

def persistent_smc_kernel(**kwargs):
    """Build a persistent SMC kernel (geometric path only)."""
    return persistent_smc(**kwargs)

adaptive_persistent_smc_kernel ¶

adaptive_persistent_smc_kernel(**kwargs)

Build an adaptive persistent SMC kernel (geometric path only).

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

def adaptive_persistent_smc_kernel(**kwargs):
    """Build an adaptive persistent SMC kernel (geometric path only)."""
    return adaptive_persistent_smc(**kwargs)

probjax.inference.adaptive_smc `module-attribute` ¶

adaptive_smc = make_smc_api(name='adaptive_smc', init_fn=init, init_params_fn=init_params, build_step_fn=build_step)

probjax.inference.persistent_smc ¶

probjax.inference.adaptive_persistent_smc ¶

probjax.inference.path_smc ¶

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."""
    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=blackjax.smc.resampling.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=blackjax.smc.resampling.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):
        import jax

        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):
        import jax

        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.likelihood_diagnostics ¶

likelihood_diagnostics(backend, key, theta, t0, data, count=None, *, num_replicates=8, batch_size=0)

Replicate likelihood estimates at one fixed theta, not across theta values.

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

def likelihood_diagnostics(
    backend, key, theta, t0, data, count=None, *, num_replicates=8, batch_size=0
):
    """Replicate likelihood estimates at one fixed theta, not across theta values."""
    if num_replicates < 2:
        raise ValueError("At least two likelihood replicates are required.")
    count = data.ts.shape[0] if count is None else count
    estimates = map_fn(
        lambda k: replay_filter(backend, k, theta, t0, data, count)[1], batch_size
    )(jax.random.split(key, num_replicates))
    return LikelihoodDiagnostics(
        estimates,
        jnp.var(estimates, ddof=1),
        jax.scipy.special.logsumexp(estimates) - jnp.log(num_replicates),
    )

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]),
    )

probjax.inference.recommend_particle_count ¶

recommend_particle_count(current_count, log_likelihood_variance, buckets, *, target_variance=1.0)

Choose a nondecreasing capacity bucket using the approximate variance ~ 1/N law.

This is a cost heuristic, not a guarantee of mixing. Call likelihood_diagnostics at fixed parameters; variance across different theta values is not estimator noise.

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

def recommend_particle_count(
    current_count, log_likelihood_variance, buckets, *, target_variance=1.0
):
    """Choose a nondecreasing capacity bucket using the approximate variance ~ 1/N law.

    This is a cost heuristic, not a guarantee of mixing. Call likelihood_diagnostics
    at fixed parameters; variance across different theta values is not estimator noise.
    """
    if target_variance <= 0 or not buckets or any(a <= 0 for a in buckets):
        raise ValueError('Supply positive buckets and a positive target variance.')
    if tuple(sorted(set(buckets))) != tuple(buckets):
        raise ValueError('Buckets must be strictly increasing.')
    sizes = jnp.asarray(buckets)
    required = current_count * jnp.maximum(
        log_likelihood_variance / target_variance, 1.0
    )
    required = jnp.where(jnp.isfinite(required), required, sizes[-1] + 1)
    index = jnp.minimum(jnp.searchsorted(sizes, required), len(buckets) - 1)
    return ParticleCountInfo(
        sizes[index], log_likelihood_variance, required > sizes[-1]
    )

probjax.inference.exchange_filter_population ¶

exchange_filter_population(key, state, new_backend, replay, *, batch_size=0)

Importance-exchange an SMC² population to a different inner particle count.

Fresh filters are replayed under q_new. The extended-target correction is L_new/L_old; replacing filters without this correction changes the target. This function is JIT-compatible for each fixed old/new capacity pair. A host controller selects a cached compiled pair when a change in array shape is needed. No proposal state or user parameter values are discarded. No resampling occurs. The normalizing correction is included in the evidence estimate.

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

def exchange_filter_population(key, state, new_backend, replay, *, batch_size=0):
    """Importance-exchange an SMC² population to a different inner particle count.

    Fresh filters are replayed under q_new. The extended-target correction is
    L_new/L_old; replacing filters without this correction changes the target.
    This function is JIT-compatible for each fixed old/new capacity pair. A host
    controller selects a cached compiled pair when a change in array shape is needed.
    No proposal state or user parameter values are discarded. No resampling occurs.
    The normalizing correction is included in the evidence estimate.
    """
    if new_backend.likelihood_kind != 'unbiased':
        raise ValueError(
            'Particle-count exchanges require an unbiased particle backend.'
        )
    n = state.log_weights.shape[0]
    if replay.ts.shape[0] == 0:
        raise ValueError('ReplayData must contain the assimilated prefix.')
    filters, logz, _ = map_fn(
        lambda args: replay_filter(
            new_backend, args[0], args[1], state.t0, replay, state.num_observations
        ),
        batch_size,
    )((jax.random.split(key, n), state.parameters))
    raw = state.log_weights + logz - state.log_likelihoods
    increment = jax.scipy.special.logsumexp(raw)
    prefix_ok = (state.num_observations <= replay.ts.shape[0]) & jnp.where(
        state.num_observations == 0,
        True,
        replay.ts[state.num_observations - 1] == state.t,
    )
    valid = jnp.isfinite(increment) & prefix_ok
    # Output has the new shape even on failure; check valid before committing.
    weights = jnp.where(valid, raw - increment, -jnp.inf)
    return state._replace(
        filter_states=filters,
        log_likelihoods=logz,
        log_weights=weights,
        log_evidence=state.log_evidence + increment,
    ), ParticleExchangeInfo(increment, jnp.where(valid, ess(weights), 0.0), valid)

probjax.inference.adapt_particle_count ¶

adapt_particle_count(key, state, backend_factory, current_count, replay, *, buckets=(64, 128, 256, 512), target_variance=1.0, num_parameters=4, num_replicates=8, batch_size=0)

Host-side automatic bucket selection; all expensive operations are compiled.

backend_factory(N) supplies fixed-capacity particle filters. Diagnostics use independent runs at a weighted sample of parameter values, then the maximum estimated log-likelihood variance across these values. Changing capacity is explicit because a single JIT carry cannot change shape. For manual compiled control use likelihood_diagnostics, recommend_particle_count, and exchange_filter_population. An exchange failure raises and never returns an apparently usable population.

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

def adapt_particle_count(
    key,
    state,
    backend_factory,
    current_count,
    replay,
    *,
    buckets=(64, 128, 256, 512),
    target_variance=1.0,
    num_parameters=4,
    num_replicates=8,
    batch_size=0,
):
    """Host-side automatic bucket selection; all expensive operations are compiled.

    backend_factory(N) supplies fixed-capacity particle filters. Diagnostics use
    independent runs at a weighted sample of parameter values, then the maximum
    estimated log-likelihood variance across these values. Changing capacity is
    explicit because a single JIT carry cannot change shape. For manual compiled
    control use likelihood_diagnostics, recommend_particle_count, and
    exchange_filter_population.
    An exchange failure raises and never returns an apparently usable population.
    """
    if current_count not in buckets or num_parameters < 1:
        raise ValueError(
            'current_count must be in buckets and num_parameters positive.'
        )
    if isinstance(state.num_observations, jax.core.Tracer):
        raise ValueError(
            'adapt_particle_count selects shapes on the host; '
            'use its compiled primitives inside JIT.'
        )
    key, select_key, diagnostic_key, exchange_key = jax.random.split(key, 4)
    indices = jax.random.categorical(
        select_key, state.log_weights, shape=(num_parameters,)
    )
    parameters = jax.tree.map(lambda x: x[indices], state.parameters)
    estimate = _diagnostic_kernel(
        backend_factory, current_count, num_replicates, batch_size
    )
    variance = estimate(
        jax.random.split(diagnostic_key, num_parameters),
        parameters,
        state.t0,
        replay,
        state.num_observations,
    ).max()
    recommendation = recommend_particle_count(
        current_count, variance, buckets, target_variance=target_variance
    )
    selected = int(recommendation.recommended_count)
    info = None
    if selected != current_count:
        state, info = _exchange_kernel(backend_factory, selected, batch_size)(
            exchange_key, state, replay
        )
        if not bool(info.valid):
            raise ValueError(
                'Particle-count exchange failed; retain the previous population.'
            )
    return AdaptiveParticleResult(state, selected, recommendation, info, key)

probjax.inference.init_streaming_window ¶

init_streaming_window(state, info, capacity)

Allocate a ring buffer using one step's info (or an abstract info prototype).

Can be carried alongside the state through arbitrary streaming chunks. The initial state must precede the first appended state. All time points must be chronological; failed outer SMC steps must not be appended.

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

def init_streaming_window(state, info, capacity):
    """Allocate a ring buffer using one step's info (or an abstract info prototype).

    Can be carried alongside the state through arbitrary streaming chunks. The
    initial state must precede the first appended state. All time points must be
    chronological; failed outer SMC steps must not be appended.
    """
    if not isinstance(capacity, int) or capacity < 1:
        raise ValueError('capacity must be a positive integer.')
    allocate = lambda x: jnp.zeros((capacity,) + x.shape, x.dtype)
    return StreamingWindow(
        state,
        jax.tree.map(allocate, state),
        jax.tree.map(allocate, info),
        jnp.zeros(capacity, dtype=state.t.dtype),
        jnp.array(0),
    )

probjax.inference.append_streaming_window ¶

append_streaming_window(window, state, info)

Append one state/diagnostic pair in O(capacity-independent) indexed writes.

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

def append_streaming_window(window, state, info):
    """Append one state/diagnostic pair in O(capacity-independent) indexed writes."""
    size = window.ts.shape[0]
    slot = window.count % size
    old = jax.tree.map(lambda x: x[slot], window.states)
    boundary = jax.lax.cond(
        window.count >= size, lambda: old, lambda: window.initial_state
    )
    return StreamingWindow(
        boundary,
        jax.tree.map(lambda x, v: x.at[slot].set(v), window.states, state),
        jax.tree.map(lambda x, v: x.at[slot].set(v), window.infos, info),
        window.ts.at[slot].set(state.t),
        window.count + 1,
    )

probjax.inference.streaming_window_trace ¶

streaming_window_trace(window)

Return (fixed-shape trace, validity mask) in chronological order.

Before capacity is filled, invalid entries are trailing zero padding; use the mask for plotting/reductions. Smoothers require an unpadded trace: on the host, slice by mask.sum(), or wait until the window is full before compiled smoothing.

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

def streaming_window_trace(window):
    """Return (fixed-shape trace, validity mask) in chronological order.

    Before capacity is filled, invalid entries are trailing zero padding; use the
    mask for plotting/reductions. Smoothers require an unpadded trace: on the host,
    slice by mask.sum(), or wait until the window is full before compiled smoothing.
    """
    size = window.ts.shape[0]
    count = jnp.minimum(window.count, size)
    start = jnp.where(window.count >= size, window.count % size, 0)
    order = (start + jnp.arange(size)) % size
    return TemporalTrace(
        window.initial_state,
        window.ts[order],
        jax.tree.map(lambda x: x[order], window.states),
        jax.tree.map(lambda x: x[order], window.infos),
    ), jnp.arange(size) < count

probjax.inference.sample_joint_paths ¶

sample_joint_paths(key, state, backend, replay, *, num_samples=100, transition_fn=None, transition_logdensity_fn=None, method='backward', batch_size=0)

Sample parameters then conditional paths; paths have shape (M,T+1,D).

ReplayData must contain exactly the assimilated prefix (no future suffix). transition_fn(theta,t0,t1)->Phi selects Gaussian backward sampling. Otherwise transition_logdensity_fn(theta,new,old,t0,t1) supplies FFBSi density. Gaussian conditional draws are exact given the SMC parameter approximation; rerun particle-filter smoothing is an additional finite-particle approximation, not an exact joint draw from the SMC² extended state. Use PGAS for further moves.

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

def sample_joint_paths(
    key,
    state,
    backend,
    replay,
    *,
    num_samples=100,
    transition_fn=None,
    transition_logdensity_fn=None,
    method='backward',
    batch_size=0,
):
    """Sample parameters then conditional paths; paths have shape (M,T+1,D).

    ReplayData must contain exactly the assimilated prefix (no future suffix).
    transition_fn(theta,t0,t1)->Phi selects Gaussian backward sampling.
    Otherwise transition_logdensity_fn(theta,new,old,t0,t1) supplies FFBSi density.
    Gaussian conditional draws are exact given the SMC parameter approximation;
    rerun particle-filter smoothing is an additional finite-particle approximation,
    not an exact joint draw from the SMC² extended state. Use PGAS for further moves.
    """
    if num_samples < 1:
        raise ValueError('num_samples must be positive.')
    if (
        not isinstance(state.num_observations, jax.core.Tracer)
        and int(state.num_observations) != replay.ts.shape[0]
    ):
        raise ValueError('Joint sampling requires exactly the assimilated data prefix.')
    if (
        transition_fn is None
        and method == 'backward'
        and transition_logdensity_fn is None
    ):
        raise ValueError(
            'Supply a Gaussian transition matrix or a particle transition density.'
        )
    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
        init_key, filter_key, sample_key = jax.random.split(key, 3)
        result = run_temporal_filter(
            backend,
            filter_key,
            backend.init(init_key, theta, state.t0),
            theta,
            replay.ts,
            replay.observations,
            observed=replay.observed,
        )
        if transition_fn is not None:
            return sample_gaussian_paths(
                sample_key, result.trace, lambda a, b: transition_fn(theta, a, b)
            )[:, 0]
        density = (
            None
            if transition_logdensity_fn is None
            else lambda new, old, a, b: transition_logdensity_fn(theta, new, old, a, b)
        )
        return sample_particle_paths(
            sample_key, result.trace, method=method, transition_logdensity_fn=density
        )[:, 0]

    paths = map_fn(draw, batch_size)((
        jax.random.split(path_key, num_samples),
        parameters,
    ))
    valid = state.num_observations == replay.ts.shape[0]
    valid = valid & (state.t == (replay.ts[-1] if replay.ts.shape[0] else state.t0))
    return JointPathSamples(
        parameters, jnp.where(valid, paths, jnp.nan), indices, valid
    )

Temporal SMC and trajectory inference¶

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

probjax.inference.kalman_backend ¶

kalman_backend(initial_fn, transition_fn, observation_fn)

Adapt the dense Kalman filter for exact linear-Gaussian likelihoods.

initial_fn(theta, t0) -> (mean, covariance) transition_fn(theta, t_previous, t_next) -> (Phi, Q) observation_fn(theta, t_next) -> (C, R)

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

def kalman_backend(initial_fn, transition_fn, observation_fn):
    """Adapt the dense Kalman filter for exact linear-Gaussian likelihoods.

    initial_fn(theta, t0) -> (mean, covariance)
    transition_fn(theta, t_previous, t_next) -> (Phi, Q)
    observation_fn(theta, t_next) -> (C, R)
    """

    def init(key, theta, t0):
        t0 = jnp.asarray(t0, dtype=jnp.result_type(t0, 0.0))
        mean, cov = initial_fn(theta, t0)
        return kalman_filter.init(mean, cov, t0)

    def step(key, state, theta, t, observation, observed=True):
        t = jnp.asarray(t, dtype=state.t.dtype)
        kernel = kalman_filter(
            lambda old, new: transition_fn(theta, old, new),
            lambda new: observation_fn(theta, new),
        )
        return jax.lax.cond(
            observed,
            lambda: kernel(state, t, observation, key),
            lambda: kernel(state, t, None, key),
        )

    return TemporalFilter(init, step, "exact")

probjax.inference.particle_backend ¶

particle_backend(initial_fn, transition_fn, log_likelihood_fn, *, ess_threshold=0.5, proposal_fn=None, proposal_logdensity_fn=None, transition_logdensity_fn=None)

Adapt the particle filter, using discrete systematic resampling.

initial_fn(key, theta, t0) -> particles (N, D), sampled from the initial law transition_fn(key, theta, particles, t_previous, t_next) -> particles log_likelihood_fn(theta, particles, observation, t_next) -> (N,)

Optional proposals use the transition sampler's signature plus observation as the final argument. Both densities take (theta, new_particles, old_particles, t_previous, t_next), with the proposal density also taking observation. Missing observations use the model transition. All three proposal arguments must be supplied together. Proposals must cover the support of the transition/observation target.

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

def particle_backend(
    initial_fn,
    transition_fn,
    log_likelihood_fn,
    *,
    ess_threshold=0.5,
    proposal_fn=None,
    proposal_logdensity_fn=None,
    transition_logdensity_fn=None,
):
    """Adapt the particle filter, using discrete systematic resampling.

    initial_fn(key, theta, t0) -> particles (N, D), sampled from the initial law
    transition_fn(key, theta, particles, t_previous, t_next) -> particles
    log_likelihood_fn(theta, particles, observation, t_next) -> (N,)

    Optional proposals use the transition sampler's signature plus ``observation``
    as the final argument. Both densities take
    (theta, new_particles, old_particles, t_previous, t_next), with the proposal
    density also taking observation. Missing observations use the model transition.
    All three proposal arguments must be supplied together. Proposals must cover
    the support of the transition/observation target.
    """
    supplied = (proposal_fn, proposal_logdensity_fn, transition_logdensity_fn)
    if any(fn is not None for fn in supplied) and not all(
        fn is not None for fn in supplied
    ):
        raise ValueError("A proposal requires its density and the transition density.")
    if not 0 <= ess_threshold <= 1:
        raise ValueError("ess_threshold must lie in [0, 1].")

    def init(key, theta, t0):
        t0 = jnp.asarray(t0, dtype=jnp.result_type(t0, 0.0))
        particles = initial_fn(key, theta, t0)
        if particles.ndim != 2 or particles.shape[0] < 1:
            raise ValueError(
                "Particle initializers must return a nonempty (N, D) array."
            )
        return ParticleFilter.init(particles, jnp.asarray(t0))

    def step(key, state, theta, t, observation, observed=True):
        t = jnp.asarray(t, dtype=state.t.dtype)
        transition = lambda key, x, t: transition_fn(key, theta, x, state.t, t)
        likelihood = lambda x, y, t: log_likelihood_fn(theta, x, y, t)
        kwargs = {}
        if proposal_fn is not None:
            kwargs = dict(
                proposal_transition_fn=lambda key, x, t: proposal_fn(
                    key, theta, x, state.t, t, observation
                ),
                proposal_logdensity_fn=lambda new, old, t: proposal_logdensity_fn(
                    theta, new, old, state.t, t, observation
                ),
                transition_logdensity_fn=lambda new, old, t: transition_logdensity_fn(
                    theta, new, old, state.t, t
                ),
            )
        kernel = ParticleFilter(
            likelihood,
            transition,
            resample_criterion=lambda ess: ess < ess_threshold,
            **kwargs,
        )
        predict = ParticleFilter(
            likelihood, transition, resample_criterion=lambda ess: ess < ess_threshold
        )
        return jax.lax.cond(
            observed,
            lambda: kernel(state, t, observation, key),
            lambda: predict(state, t, None, key),
        )

    return TemporalFilter(init, step, "unbiased")

probjax.inference.run_temporal_filter ¶

run_temporal_filter(backend, key, state, theta, ts, observations, *, observed=None, history='full')

Scan an incremental filter; history is 'full', 'none', or a window size.

The grid must increase from state.t. A positive integer keeps a bounded ring buffer; 'none' stores no time history. No growing arrays enter the scan carry. Streaming reproduces a run by repeatedly splitting key, step_key and calling backend.step. Returned log_likelihood covers only this run's intervals.

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

def run_temporal_filter(
    backend, key, state, theta, ts, observations, *, observed=None, history="full"
):
    """Scan an incremental filter; ``history`` is 'full', 'none', or a window size.

    The grid must increase from state.t. A positive integer keeps a bounded ring
    buffer; 'none' stores no time history. No growing arrays enter the scan carry.
    Streaming reproduces a run by repeatedly splitting ``key, step_key`` and
    calling backend.step. Returned log_likelihood covers only this run's intervals.
    """
    ts, observations, observed = _observations(ts, observations, observed)
    n = ts.shape[0]
    if history not in ("full", "none") and (
        not isinstance(history, int) or isinstance(history, bool) or history < 1
    ):
        raise ValueError("history must be 'full', 'none', or a positive integer.")

    def advance(carry, data):
        key, state, logz = carry
        key, step_key = jax.random.split(key)
        t, y, mask = data
        state, info = backend.step(step_key, state, theta, t, y, mask)
        return (key, state, logz + info.log_likelihood), (state, info)

    zero = jnp.asarray(0.0, dtype=jnp.result_type(jax.tree.leaves(state)[0], 0.0))
    carry = (key, state, zero)
    if history == "none":

        def discard(carry, data):
            carry, _ = advance(carry, data)
            return carry, None

        (key, final, logz), _ = jax.lax.scan(
            discard, carry, (ts, observations, observed)
        )
        return TemporalResult(final, logz, key, None)
    if history == "full" or n == 0:
        (key, final, logz), (states, infos) = jax.lax.scan(
            advance, carry, (ts, observations, observed)
        )
        return TemporalResult(final, logz, key, TemporalTrace(state, ts, states, infos))

    capacity = min(history, n)
    info_shape = jax.eval_shape(
        backend.step, key, state, theta, ts[0], observations[0], observed[0]
    )[1]
    allocate = lambda x: jnp.zeros((capacity,) + x.shape, x.dtype)
    states = jax.tree.map(allocate, state)
    infos = jax.tree.map(allocate, info_shape)

    def window_step(carry, data):
        running, boundary, states, infos = carry
        i, t, y, mask = data
        slot = i % capacity
        old = jax.tree.map(lambda x: x[slot], states)
        boundary = jax.lax.cond(i >= capacity, lambda: old, lambda: boundary)
        running, (next_state, info) = advance(running, (t, y, mask))
        states = jax.tree.map(lambda x, v: x.at[slot].set(v), states, next_state)
        infos = jax.tree.map(lambda x, v: x.at[slot].set(v), infos, info)
        return (running, boundary, states, infos), None

    ((key, final, logz), boundary, states, infos), _ = jax.lax.scan(
        window_step,
        (carry, state, states, infos),
        (jnp.arange(n), ts, observations, observed),
    )
    order = (jnp.arange(capacity) + n) % capacity
    states, infos = jax.tree.map(lambda x: x[order], (states, infos))
    return TemporalResult(
        final, logz, key, TemporalTrace(boundary, ts[-capacity:], states, infos)
    )

probjax.inference.temporal_smc ¶

temporal_smc(backend, logprior_fn, *, proposal_fn=None, num_rejuvenation_steps=1, ess_threshold=0.5, batch_size=0, adaptive_proposal=False, proposal_scale=1.0, adaptive_num_steps=False, max_rejuvenation_steps=16, target_accepted_moves=2.0, target_acceptance=0.234, tempering_ess=None, max_tempering_steps=64, mcmc_kernel=None, mcmc_parameters=None, mcmc_kernel_kwargs=None)

Build temporal SMC using exact or unbiased incremental likelihoods.

Initial particles must be equally weighted prior draws. proposal_fn(key,theta) returns (candidate, log_q_reverse_minus_forward). adaptive_proposal=True instead uses a Gaussian random walk with weighted population covariance and an acceptance-tuned scale, frozen during each rejuvenation sweep.

Exact backends may instead supply an existing probjax mcmc_kernel and its unbatched parameter dictionary (e.g. HMC). This uses differentiable prefix replay. No gradient-based move is applied to noisy particle likelihoods.

tempering_ess optionally bridges each observation with conditional ESS steps. For stochastic backends these are EXTENDED-SPACE targets with retained filter likelihood estimates, not powers of the marginalized observation likelihood. The beta=1 endpoint is ordinary SMC². Prefix replay retains both L_previous and L_current so pseudo-marginal proposals preserve every bridge target.

Each step is atomic: invalid replay/time or an unfinished bridge returns the input state with info.valid=False; increase max_tempering_steps and retry. All-impossible observations are reported as invalid (never silently uniform). batch_size bounds population mapping memory; zero uses full vectorization.

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

def temporal_smc(
    backend,
    logprior_fn,
    *,
    proposal_fn=None,
    num_rejuvenation_steps=1,
    ess_threshold=0.5,
    batch_size=0,
    adaptive_proposal=False,
    proposal_scale=1.0,
    adaptive_num_steps=False,
    max_rejuvenation_steps=16,
    target_accepted_moves=2.0,
    target_acceptance=0.234,
    tempering_ess=None,
    max_tempering_steps=64,
    mcmc_kernel=None,
    mcmc_parameters=None,
    mcmc_kernel_kwargs=None,
):
    """Build temporal SMC using exact or unbiased incremental likelihoods.

    Initial particles must be equally weighted prior draws. proposal_fn(key,theta)
    returns (candidate, log_q_reverse_minus_forward). adaptive_proposal=True
    instead uses a Gaussian random walk with weighted population covariance and
    an acceptance-tuned scale, frozen during each rejuvenation sweep.

    Exact backends may instead supply an existing probjax mcmc_kernel and its
    unbatched parameter dictionary (e.g. HMC). This uses differentiable prefix
    replay. No gradient-based move is applied to noisy particle likelihoods.

    tempering_ess optionally bridges each observation with conditional ESS steps.
    For stochastic backends these are EXTENDED-SPACE targets with retained filter
    likelihood estimates, not powers of the marginalized observation likelihood.
    The beta=1 endpoint is ordinary SMC². Prefix replay retains both L_previous
    and L_current so pseudo-marginal proposals preserve every bridge target.

    Each step is atomic: invalid replay/time or an unfinished bridge returns the
    input state with info.valid=False; increase max_tempering_steps and retry.
    All-impossible observations are reported as invalid (never silently uniform).
    batch_size bounds population mapping memory; zero uses full vectorization.
    """
    if backend.likelihood_kind not in ('exact', 'unbiased'):
        raise ValueError(
            'Parameter SMC requires an exact or unbiased likelihood backend.'
        )
    if not 0 <= ess_threshold <= 1:
        raise ValueError('ess_threshold must lie in [0, 1].')
    if not isinstance(num_rejuvenation_steps, int) or num_rejuvenation_steps < 0:
        raise ValueError('num_rejuvenation_steps must be a nonnegative integer.')
    if adaptive_num_steps and (
        not isinstance(max_rejuvenation_steps, int)
        or max_rejuvenation_steps < num_rejuvenation_steps
        or target_accepted_moves <= 0
    ):
        raise ValueError('Invalid adaptive move-count budget.')
    if tempering_ess is not None and not 0 < tempering_ess < 1:
        raise ValueError('tempering_ess must lie strictly between zero and one.')
    if not isinstance(max_tempering_steps, int) or max_tempering_steps < 1:
        raise ValueError('max_tempering_steps must be positive.')
    if proposal_scale <= 0 or not 0 < target_acceptance < 1:
        raise ValueError('proposal_scale and target_acceptance must be positive/valid.')
    if sum((proposal_fn is not None, adaptive_proposal, mcmc_kernel is not None)) > 1:
        raise ValueError(
            'Choose one proposal function, adaptive Gaussian, or MCMC kernel.'
        )
    if mcmc_kernel is not None:
        if backend.likelihood_kind != 'exact':
            raise ValueError(
                'Existing MCMC kernels require an exact likelihood backend.'
            )
        if mcmc_parameters is None:
            raise ValueError('mcmc_parameters must be supplied with an MCMC kernel.')
        mcmc_init, mcmc_step = make_mcmc_adapter(
            mcmc_kernel, **(mcmc_kernel_kwargs or {})
        )
    moves = (
        proposal_fn is not None or adaptive_proposal or mcmc_kernel is not None
    ) and num_rejuvenation_steps > 0

    if adaptive_num_steps and not moves:
        raise ValueError(
            'Adaptive move counts require an enabled rejuvenation proposal.'
        )

    def init(key, parameters, t0):
        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))
        states = map_fn(lambda args: backend.init(args[0], args[1], t0), batch_size)((
            jax.random.split(key, n),
            parameters,
        ))
        priors = map_fn(logprior_fn, batch_size)(parameters)
        if priors.shape != (n,):
            raise ValueError('logprior_fn must return a scalar per parameter.')
        zeros = jnp.zeros_like(priors)
        return TemporalSMCState(
            parameters,
            states,
            zeros - jnp.log(n),
            zeros,
            jnp.zeros((), priors.dtype),
            t0,
            t0,
            jnp.array(0, jnp.int32),
            jnp.asarray(proposal_scale),
            jnp.asarray(num_rejuvenation_steps, dtype=jnp.int32),
            jnp.arange(n, dtype=jnp.int32),
        )

    def step(key, state, t, observation, observed=True, replay=None):
        t = jnp.asarray(t, state.t.dtype)
        count = state.num_observations + 1
        replay_valid = jnp.array(True)
        if moves:
            if replay is None:
                raise ValueError(
                    'Rejuvenation requires ReplayData for the full observed prefix.'
                )
            if replay.ts.ndim != 1 or replay.ts.shape[0] == 0:
                raise ValueError(
                    'ReplayData cannot be empty when assimilating observations.'
                )
            if (
                replay.observations.shape[0] != replay.ts.shape[0]
                or replay.observed.shape != replay.ts.shape
            ):
                raise ValueError('ReplayData arrays must share their time dimension.')
            replay_valid = (count <= replay.ts.shape[0]) & (replay.ts[count - 1] == t)
            if not isinstance(replay_valid, jax.core.Tracer) and not bool(replay_valid):
                raise ValueError(
                    'ReplayData must contain the full prefix through this time.'
                )
        n = state.log_weights.shape[0]
        key, predict_key = jax.random.split(key)
        states, infos = map_fn(
            lambda args: backend.step(
                args[0], args[1], args[2], t, observation, observed
            ),
            batch_size,
        )((jax.random.split(predict_key, n), state.filter_states, state.parameters))
        likelihoods = state.log_likelihoods + infos.log_likelihood

        def rejuvenate(
            key, parameters, states, likelihoods, previous_ll, beta, covariance, scale
        ):
            def one(args):
                key, theta, current, ll, previous = args

                def target(p):
                    _, new, old = replay_filter(
                        backend, key, p, state.t0, replay, count, differentiable=True
                    )
                    return logprior_fn(p) + old + beta * (new - old)

                def move(_, carry):
                    key, theta, current, ll, previous, accepted = carry
                    key, proposal_key, replay_key, accept_key = jax.random.split(key, 4)
                    if mcmc_kernel is not None:
                        init_key, kernel_key = jax.random.split(proposal_key)
                        chain = mcmc_init(theta, target, rng_key=init_key)
                        chain, info = mcmc_step(
                            kernel_key, chain, target, **mcmc_parameters
                        )
                        theta = chain.position
                        current, ll, previous = replay_filter(
                            backend, replay_key, theta, state.t0, replay, count
                        )
                        accepted += info.acceptance_rate
                    else:
                        if adaptive_proposal:
                            flat, unravel = ravel_pytree(theta)
                            noise = jax.random.normal(proposal_key, flat.shape)
                            proposed = unravel(
                                flat
                                + scale
                                * (2.38 / jnp.sqrt(flat.size))
                                * (jnp.linalg.cholesky(covariance) @ noise)
                            )
                            correction = 0.0
                        else:
                            proposed, correction = proposal_fn(proposal_key, theta)
                        candidate, new_ll, old_ll = replay_filter(
                            backend, replay_key, proposed, state.t0, replay, count
                        )
                        ratio = (
                            logprior_fn(proposed)
                            - logprior_fn(theta)
                            + old_ll
                            - previous
                            + beta * ((new_ll - old_ll) - (ll - previous))
                            + correction
                        )
                        (theta, current, ll, previous), (accept, _, _) = (
                            static_binomial_sampling(
                                accept_key,
                                jnp.where(jnp.isnan(ratio), -jnp.inf, ratio),
                                (theta, current, ll, previous),
                                (proposed, candidate, new_ll, old_ll),
                            )
                        )
                        accepted += accept
                    return key, theta, current, ll, previous, accepted

                _, theta, current, ll, previous, accepted = jax.lax.fori_loop(
                    0,
                    state.move_steps,
                    move,
                    (key, theta, current, ll, previous, jnp.array(0.0)),
                )
                return theta, current, ll, previous, accepted / state.move_steps

            parameters, states, ll, old, rates = map_fn(one, batch_size)((
                jax.random.split(key, n),
                parameters,
                states,
                likelihoods,
                previous_ll,
            ))
            return parameters, states, ll, old, rates.mean()

        def stage(carry):
            (
                key,
                parameters,
                states,
                ll,
                old,
                weights,
                logz,
                beta,
                stages,
                accepted,
                any_resampled,
                parents,
                scale,
                valid,
                last_ess,
            ) = carry
            key, resample_key, move_key = jax.random.split(key, 3)
            increments = ll - old
            if tempering_ess is None:
                delta = 1.0 - beta
            else:

                def objective(delta):
                    power = jnp.where(delta == 0.0, 0.0, delta * increments)
                    log_cess = (
                        jnp.log(n)
                        + 2 * jax.scipy.special.logsumexp(weights + power)
                        - jax.scipy.special.logsumexp(weights + 2 * power)
                    )
                    return log_cess - jnp.log(n * tempering_ess)

                delta = dichotomy(objective, jnp.array(0.0), 1.0 - beta)
            beta_next = jnp.minimum(beta + delta, 1.0)
            raw = weights + delta * increments
            increment = jax.scipy.special.logsumexp(raw)
            valid = valid & jnp.isfinite(increment) & jnp.isfinite(delta) & (delta > 0)
            weights = jnp.where(valid, raw - increment, -jnp.inf)
            last_ess = jnp.where(valid, effective_sample_size(weights), 0.0)
            resampled = valid & ((last_ess < ess_threshold * n) | (beta_next < 1.0))
            covariance = (
                population_covariance(parameters, weights)
                if adaptive_proposal
                else jnp.zeros((1, 1))
            )
            idx = jax.lax.cond(
                resampled,
                lambda: systematic(resample_key, jnp.exp(weights), n),
                lambda: jnp.arange(n, dtype=jnp.int32),
            )
            parameters, states, ll, old, parents = jax.tree.map(
                lambda x: x[idx], (parameters, states, ll, old, parents)
            )
            weights = jnp.where(resampled, -jnp.log(n), weights)
            rate = jnp.array(0.0)
            if moves:
                parameters, states, ll, old, rate = jax.lax.cond(
                    resampled,
                    lambda: rejuvenate(
                        move_key,
                        parameters,
                        states,
                        ll,
                        old,
                        beta_next,
                        covariance,
                        scale,
                    ),
                    lambda: (parameters, states, ll, old, rate),
                )
            if adaptive_proposal:
                scale = jnp.where(
                    resampled,
                    jnp.clip(
                        update_scale_from_acceptance_rate(
                            scale, rate, target_acceptance
                        ),
                        1e-3,
                        1e3,
                    ),
                    scale,
                )
            return (
                key,
                parameters,
                states,
                ll,
                old,
                weights,
                logz + increment,
                beta_next,
                stages + 1,
                accepted + rate,
                any_resampled | resampled,
                parents,
                scale,
                valid,
                last_ess,
            )

        carry = (
            key,
            state.parameters,
            states,
            likelihoods,
            state.log_likelihoods,
            state.log_weights,
            jnp.zeros_like(state.log_evidence),
            jnp.array(0.0),
            jnp.array(0),
            jnp.array(0.0),
            jnp.array(False),
            jnp.arange(n, dtype=jnp.int32),
            state.proposal_scale,
            replay_valid & (t > state.t),
            jnp.array(0.0),
        )
        # A single untempered step avoids a while-loop around the common path.
        if tempering_ess is None:
            carry = stage(carry)
        else:
            carry = jax.lax.while_loop(
                lambda c: (c[7] < 1.0) & (c[8] < max_tempering_steps) & c[13],
                stage,
                carry,
            )
        (
            _,
            parameters,
            states,
            ll,
            _,
            weights,
            increment,
            beta,
            stages,
            acceptance,
            resampled,
            parents,
            scale,
            valid,
            last_ess,
        ) = carry
        valid = valid & (beta >= 1.0)
        next_state = TemporalSMCState(
            parameters,
            states,
            weights,
            ll,
            state.log_evidence + increment,
            state.t0,
            t,
            count,
            scale,
            jnp.where(
                adaptive_num_steps & resampled,
                jnp.clip(
                    jnp.ceil(
                        target_accepted_moves
                        / jnp.maximum(acceptance / jnp.maximum(stages, 1), 1e-6)
                    ),
                    1,
                    max_rejuvenation_steps,
                ).astype(jnp.int32),
                state.move_steps,
            ),
            state.lineages[parents],
        )
        next_state = jax.lax.cond(valid, lambda: next_state, lambda: state)
        return next_state, TemporalSMCInfo(
            jnp.where(valid, increment, 0.0),
            last_ess,
            resampled,
            parents,
            acceptance / jnp.maximum(stages, 1),
            valid,
            stages,
            beta,
        )

    return TemporalSMC(init, step, moves)

probjax.inference.run_temporal_smc ¶

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

Run parameter SMC; return TemporalResult with optional population history.

For a fresh state, replay defaults to this run's data. When resuming a stream with rejuvenation, pass ReplayData starting at state.t0, not just the new chunk. History='none' stores only the final population; integer windows are bounded.

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

def run_temporal_smc(
    kernel, key, state, ts, observations, *, observed=None, history="none", replay=None
):
    """Run parameter SMC; return TemporalResult with optional population history.

    For a fresh state, replay defaults to this run's data. When resuming a stream
    with rejuvenation, pass ReplayData starting at state.t0, not just the new chunk.
    History='none' stores only the final population; integer windows are bounded.
    """
    from probjax.inference.filtering.temporal import (
        TemporalFilter,
        _observations,
        run_temporal_filter,
    )

    ts, observations, observed = _observations(ts, observations, observed)
    if replay is None:
        if (
            kernel.requires_replay
            and not isinstance(state.num_observations, jax.core.Tracer)
            and int(state.num_observations) != 0
        ):
            raise ValueError(
                "Resuming with rejuvenation requires full-prefix ReplayData."
            )
        replay = ReplayData(ts, observations, observed)
    else:
        replay = ReplayData(*_observations(*replay))
    adapter = TemporalFilter(
        kernel.init,
        lambda key, state, theta, t, y, mask: kernel.step(
            key, state, t, y, mask, replay=replay
        ),
        "exact",  # runner does not inspect this; this is an outer evidence increment.
    )
    return run_temporal_filter(
        adapter, key, state, None, ts, observations, observed=observed, history=history
    )

probjax.inference.sample_gaussian_paths ¶

sample_gaussian_paths(key, trace, transition_fn, *, num_samples=1)

Sample joint Gaussian trajectories, preserving cross-time dependence.

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

def sample_gaussian_paths(key, trace, transition_fn, *, num_samples=1):
    """Sample joint Gaussian trajectories, preserving cross-time dependence."""
    if num_samples < 1:
        raise ValueError("num_samples must be positive.")
    states, gains = _gaussian_backward_data(trace, transition_fn)
    key, final_key = jax.random.split(key)
    final = _normal(final_key, states.mean[-1], states.cov[-1], num_samples)

    def backward(carry, data):
        key, next_x = carry
        key, draw_key = jax.random.split(key)
        mean, cov, pred_mean, pred_cov, gain = data
        conditional_mean = mean + (next_x - pred_mean) @ gain.T
        conditional_cov = cov - gain @ pred_cov @ gain.T
        x = _normal(draw_key, conditional_mean, conditional_cov, num_samples)
        return (key, x), x

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

probjax.inference.smooth_gaussian_path ¶

smooth_gaussian_path(trace, transition_fn)

RTS means/covariances including the initial state; transition_fn(t0,t1)->Phi.

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

def smooth_gaussian_path(trace, transition_fn):
    """RTS means/covariances including the initial state; transition_fn(t0,t1)->Phi."""
    states, gains = _gaussian_backward_data(trace, transition_fn)

    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_particle_paths ¶

sample_particle_paths(key, trace, *, num_samples=1, method='backward', transition_logdensity_fn=None)

Return joint paths (time INCLUDING initial state, samples, dimension).

'ancestry' traces stored discrete parents in O(TM); 'backward' uses FFBSi in O(TN*M), requiring log p(x_next|x_previous, t_previous, t_next). Traces must come from discrete resampling, as in particle_backend.

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

def sample_particle_paths(
    key, trace, *, num_samples=1, method="backward", transition_logdensity_fn=None
):
    """Return joint paths (time INCLUDING initial state, samples, dimension).

    'ancestry' traces stored discrete parents in O(T*M); 'backward' uses FFBSi
    in O(T*N*M), requiring log p(x_next|x_previous, t_previous, t_next).
    Traces must come from discrete resampling, as in particle_backend.
    """
    if num_samples < 1:
        raise ValueError("num_samples must be positive.")
    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 transition_logdensity_fn 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,
        transition_logdensity_fn,
        num_samples=num_samples,
    )
    return paths

probjax.inference.particle_gibbs ¶

particle_gibbs(key, reference_path, theta, t0, ts, observations, *, initial_fn, transition_fn, transition_logdensity_fn, log_likelihood_fn, num_particles=32, observed=None, ancestor_sampling=True)

One bootstrap conditional-SMC/PGAS update, returning a new joint path.

reference_path has shape (len(ts)+1, D), including x(t0). Callback signatures match particle_backend; the transition density takes scalar states and returns a scalar: (theta, x_next, x_previous, t_previous, t_next). initial_fn samples exactly num_particles independent draws from the initial law. Observation likelihoods are batched. Multinomial resampling at every step is intentional: conditioning ordinary systematic resampling by pinning one index is invalid.

The kernel preserves the full smoothing target for a 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,
    theta,
    t0,
    ts,
    observations,
    *,
    initial_fn,
    transition_fn,
    transition_logdensity_fn,
    log_likelihood_fn,
    num_particles=32,
    observed=None,
    ancestor_sampling=True,
):
    """One bootstrap conditional-SMC/PGAS update, returning a new joint path.

    reference_path has shape (len(ts)+1, D), including x(t0). Callback signatures
    match particle_backend; the transition density takes scalar states and returns
    a scalar: (theta, x_next, x_previous, t_previous, t_next). initial_fn samples
    exactly num_particles independent draws from the initial law. Observation
    likelihoods are batched. Multinomial resampling at every step is intentional:
    conditioning ordinary systematic resampling by pinning one index is invalid.

    The kernel preserves the full smoothing target for a fixed theta. Repeated
    calls form an MCMC chain; one call is not an independent posterior draw.
    """
    ts, observations, observed = _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).")
    key, initial_key = jax.random.split(key)
    particles = initial_fn(initial_key, theta, t0)
    if particles.shape != (num_particles, reference_path.shape[1]):
        raise ValueError("initial_fn must return (num_particles, path_dimension).")
    particles = particles.at[0].set(reference_path[0])
    initial_particles = particles
    weights = jnp.full((num_particles,), -jnp.log(num_particles))

    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: transition_logdensity_fn(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)
        particles = transition_fn(predict_key, theta, particles[parents], previous, t)
        particles = particles.at[0].set(reference)
        weights = jax.lax.cond(
            mask,
            lambda: log_likelihood_fn(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, mask).

step returns (state, info); info.log_likelihood is the incremental log likelihood. likelihood_kind describes the likelihood, not its logarithm: 'exact', 'unbiased', or 'approximate'. SMC² requires an unbiased likelihood.

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

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

    ``step`` returns (state, info); info.log_likelihood is the incremental log
    likelihood. ``likelihood_kind`` describes the likelihood, not its logarithm:
    'exact', 'unbiased', or 'approximate'. SMC² requires an unbiased likelihood.
    """

    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/temporal.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/temporal.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
    requires_replay: bool = False

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
    move_steps: Any
    lineages: Any

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
    num_tempering_steps: Any
    tempering_param: Any

probjax.inference.ReplayData ¶

Bases: NamedTuple

Complete data prefix starting at t0; unused future observations are permitted.

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

class ReplayData(NamedTuple):
    """Complete data prefix starting at t0; unused future observations are permitted."""

    ts: Any
    observations: Any
    observed: Any

Filtering and smoothing¶

probjax.inference.kalman_filter ¶

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

    @staticmethod
    def default_unpack(state, info):
        return (state.mean, state.cov)

default_solve ¶

default_solve(S, res, dense_mem_limit=200)

Solve S @ x = res for the Kalman gain and residual.

Chooses between dense factorization and batched PCG based on total memory required. Dense needs to materialize S (obs_dim^2) on top of the already- materialized res (nrhs * obs_dim). PCG avoids materializing S entirely.

Decision rule: use dense when obs_dim^2 * 8 bytes < dense_mem_limit MB, i.e. when the materialized S matrix fits comfortably in memory. This naturally accounts for the RHS shape: for small obs_dim, dense factorization is O(obs_dim^3) and amortizes over all nrhs columns cheaply. For large obs_dim, PCG avoids the O(obs_dim^3) factorization and O(obs_dim^2) memory.

Parameters:

Name	Description	Default
`S`	SPD system matrix (obs_dim x obs_dim), array or LinearOperator.	required
`res`	RHS matrix, shape (nrhs, obs_dim).	required
`dense_mem_limit`	Max MB for the materialized S matrix. Default 200 MB, corresponding to obs_dim ~ 5000 in f64 or ~7000 in f32.	`200`

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

def default_solve(S, res, dense_mem_limit=200):
    """Solve S @ x = res for the Kalman gain and residual.

    Chooses between dense factorization and batched PCG based on total memory
    required. Dense needs to materialize S (obs_dim^2) on top of the already-
    materialized res (nrhs * obs_dim). PCG avoids materializing S entirely.

    Decision rule: use dense when obs_dim^2 * 8 bytes < dense_mem_limit MB,
    i.e. when the materialized S matrix fits comfortably in memory. This
    naturally accounts for the RHS shape: for small obs_dim, dense factorization
    is O(obs_dim^3) and amortizes over all nrhs columns cheaply. For large
    obs_dim, PCG avoids the O(obs_dim^3) factorization and O(obs_dim^2) memory.

    Args:
        S: SPD system matrix (obs_dim x obs_dim), array or LinearOperator.
        res: RHS matrix, shape (nrhs, obs_dim).
        dense_mem_limit: Max MB for the materialized S matrix. Default 200 MB,
            corresponding to obs_dim ~ 5000 in f64 or ~7000 in f32.
    """
    if isinstance(S, LinearOperator):
        obs_dim = S.out_dim
    else:
        S = jnp.asarray(S)
        obs_dim = S.shape[0]

    res = jnp.asarray(res)

    # Estimate S memory in MB (use f64 = 8 bytes as upper bound)
    s_mem_mb = obs_dim * obs_dim * 8 / (1024 * 1024)

    if isinstance(S, LinearOperator) and s_mem_mb > dense_mem_limit:
        # Batched PCG — no materialization
        matvec = S.operator
        rhs = res.T  # (obs_dim, nrhs)
        if rhs.ndim == 1:
            return jax.scipy.sparse.linalg.cg(matvec, rhs, tol=1e-4)[0]
        else:
            X, _info = batched_pcg_solve(matvec, rhs, tol=1e-4, block_size=128)
            return X.T
    else:
        # Dense solve — materialize S if needed, then factorize once
        if isinstance(S, LinearOperator):
            S = S.as_array()
        return jax.scipy.linalg.solve(S, res.T, assume_a="pos").T

default_logdet ¶

default_logdet(S, dense_mem_limit=200)

Compute the logdet with Lanczos for large operators, otherwise slogdet.

Uses the same memory-based decision as default_solve: if materializing S would exceed dense_mem_limit MB, use matrix-free Lanczos SLQ instead.

Benchmarks (CPU, f64): dense slogdet (including materialization) is faster until dim ~2000. Above that, Lanczos avoids O(n^2) materialization and O(n^3) factorization.

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

def default_logdet(S, dense_mem_limit=200):
    """Compute the logdet with Lanczos for large operators, otherwise slogdet.

    Uses the same memory-based decision as default_solve: if materializing S
    would exceed dense_mem_limit MB, use matrix-free Lanczos SLQ instead.

    Benchmarks (CPU, f64): dense slogdet (including materialization) is faster
    until dim ~2000. Above that, Lanczos avoids O(n^2) materialization and
    O(n^3) factorization.
    """
    if isinstance(S, LinearOperator):
        dim = S.out_dim
        s_mem_mb = dim * dim * 8 / (1024 * 1024)
        if s_mem_mb > dense_mem_limit:
            return lanczos_logdet(S, num_steps=min(500, dim))
        S = S.as_array()
    return jnp.linalg.slogdet(jnp.asarray(S)).logabsdet

probjax.inference.extended_kalman_filter ¶

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 = init
    build_kernel = build_kernel

    @staticmethod
    def default_unpack(state, info):
        return (state.mean, state.cov)

build_kernel ¶

build_kernel(transition_model_fn, observation_model_fn, linear_solve=None, logdet_fn=None)

Build an Extended Kalman filter kernel.

Parameters:

Name	Type	Description	Default
`transition_model_fn`	`Callable`	(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`	`Callable`	(x, cov, t) -> (y_pred, C, R) Returns the nonlinear predicted observation, the Jacobian of the observation function, and the observation noise covariance.	required
`linear_solve`	`Optional[Callable]`	Optional custom solve function (see kalman_filter).	`None`
`logdet_fn`	`Optional[Callable]`	Optional custom logdet function (see kalman_filter).	`None`

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

def build_kernel(
    transition_model_fn: Callable,
    observation_model_fn: Callable,
    linear_solve: Optional[Callable] = None,
    logdet_fn: Optional[Callable] = None,
) -> Callable:
    """Build an Extended Kalman filter kernel.

    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.
        linear_solve: Optional custom solve function (see kalman_filter).
        logdet_fn: Optional custom logdet function (see kalman_filter).
    """

    def kernel(
        state: KalmanFilterState,
        t: Optional[ArrayLike] = None,
        observed: Optional[ArrayLike] = None,
        rng_key: Optional[jnp.ndarray] = None,
    ) -> Tuple[KalmanFilterState, KalmanFilterInfo]:
        mu0 = state.mean
        cov0 = state.cov
        t_old = state.t
        is_observed = observed is not None

        # Predict: nonlinear mean + linearized covariance
        mu1_, Phi, Q = transition_model_fn(mu0, cov0, t_old, t)
        cov1_ = Phi @ cov0 @ Phi.T + Q
        # Materialize predicted covariance — needed for update step
        if isinstance(cov1_, LinearOperator):
            cov1_ = cov1_.as_array()

        if is_observed:
            y_, C, R = observation_model_fn(mu1_, cov1_, t)

            solve = default_solve if linear_solve is None else linear_solve
            _logdet_fn = logdet_fn if logdet_fn is not None else default_logdet

            mu1, cov1, log_likelihood = _kalman_update(
                mu1_, cov1_, y_, observed, C, R, solve, _logdet_fn
            )

            # Ensure symmetry (EKF linearization can introduce asymmetry)
            cov1 = 0.5 * (cov1 + cov1.T)

            return KalmanFilterState(mu1, cov1, t), KalmanFilterInfo(
                mu1_, cov1_, log_likelihood
            )
        else:
            return KalmanFilterState(mu1_, cov1_, t), KalmanFilterInfo(
                mu1_, cov1_, jnp.array(0.0)
            )

    return kernel

make_linearized_transition ¶

make_linearized_transition(transition_fn, Q_fn, in_dim, out_dim=None, materialize=False)

Build an EKF transition model.

By default the Jacobian Phi = df/dx is returned as a matrix-free LinearOperator. Set materialize=True to return a dense array instead (faster for small state dimensions, but O(n^2) memory).

Parameters:

Name	Description	Default
`transition_fn`	f(x, t_old, t) -> x_new. Nonlinear transition.	required
`Q_fn`	Either a callable (t_old, t) -> Q returning the process noise covariance, or a fixed array / LinearOperator.	required
`in_dim`	State dimension.	required
`out_dim`	Output dimension. Defaults to in_dim.	`None`
`materialize`	If True, return Phi as a dense array via jax.jacfwd. If False (default), return Phi as a matrix-free LinearOperator.	`False`

Returns:

Name	Type	Description
`transition_model_fn`		(x, cov, t_old, t) -> (x_pred, Phi, Q)

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

def make_linearized_transition(
    transition_fn, Q_fn, in_dim, out_dim=None, materialize=False
):
    """Build an EKF transition model.

    By default the Jacobian Phi = df/dx is returned as a matrix-free
    LinearOperator. Set materialize=True to return a dense array instead
    (faster for small state dimensions, but O(n^2) memory).

    Args:
        transition_fn: f(x, t_old, t) -> x_new.  Nonlinear transition.
        Q_fn: Either a callable (t_old, t) -> Q returning the process noise
            covariance, or a fixed array / LinearOperator.
        in_dim: State dimension.
        out_dim: Output dimension. Defaults to in_dim.
        materialize: If True, return Phi as a dense array via jax.jacfwd.
            If False (default), return Phi as a matrix-free LinearOperator.

    Returns:
        transition_model_fn: (x, cov, t_old, t) -> (x_pred, Phi, Q)
    """
    if out_dim is None:
        out_dim = in_dim

    def transition_model_fn(x, cov, t_old, t):
        x_pred = transition_fn(x, t_old, t)
        if materialize:
            Phi = jax.jacfwd(transition_fn)(x, t_old, t)
        else:
            Phi = _jacobian_as_linear_operator(
                lambda x: transition_fn(x, t_old, t),
                x,
                in_dim=in_dim,
                out_dim=out_dim,
            )
        Q = Q_fn(t_old, t) if callable(Q_fn) else Q_fn
        return x_pred, Phi, Q

    return transition_model_fn

make_linearized_observation ¶

make_linearized_observation(observation_fn, R_fn, in_dim, out_dim, materialize=False)

Build an EKF observation model.

By default the Jacobian C = dh/dx is returned as a matrix-free LinearOperator. Set materialize=True to return a dense array instead.

Parameters:

Name	Description	Default
`observation_fn`	h(x, t) -> y. Nonlinear observation.	required
`R_fn`	Either a callable (t,) -> R returning the observation noise covariance, or a fixed array / LinearOperator.	required
`in_dim`	State (input) dimension.	required
`out_dim`	Observation (output) dimension.	required
`materialize`	If True, return C as a dense array via jax.jacfwd. If False (default), return C as a matrix-free LinearOperator.	`False`

Returns:

Name	Type	Description
`observation_model_fn`		(x, cov, t) -> (y_pred, C, R)

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

def make_linearized_observation(
    observation_fn, R_fn, in_dim, out_dim, materialize=False
):
    """Build an EKF observation model.

    By default the Jacobian C = dh/dx is returned as a matrix-free
    LinearOperator. Set materialize=True to return a dense array instead.

    Args:
        observation_fn: h(x, t) -> y.  Nonlinear observation.
        R_fn: Either a callable (t,) -> R returning the observation noise
            covariance, or a fixed array / LinearOperator.
        in_dim: State (input) dimension.
        out_dim: Observation (output) dimension.
        materialize: If True, return C as a dense array via jax.jacfwd.
            If False (default), return C as a matrix-free LinearOperator.

    Returns:
        observation_model_fn: (x, cov, t) -> (y_pred, C, R)
    """

    def observation_model_fn(x, cov, t):
        y_pred = observation_fn(x, t)
        if materialize:
            C = jax.jacfwd(observation_fn)(x, t)
        else:
            C = _jacobian_as_linear_operator(
                lambda x: observation_fn(x, t),
                x,
                in_dim=in_dim,
                out_dim=out_dim,
            )
        R = R_fn(t) if callable(R_fn) else R_fn
        return y_pred, C, R

    return observation_model_fn

make_continuous_transition ¶

make_continuous_transition(drift_fn, diffusion_matrix, in_dim, materialize=False)

Build an EKF transition model from continuous SDE.

For an SDE dx = f(x,t) dt + B dw, this computes: - x_pred = x + f(x, t_old) * (t - t_old) (Euler step) - Phi, Q via matrix_fraction_decomposition

Since matrix_fraction_decomposition must materialize the 2d x 2d block matrix for expm, Phi is always dense regardless of the materialize flag. The flag only controls whether the separate Jacobian A is kept dense or discarded (it's computed regardless for the decomposition).

Parameters:

Name	Description	Default
`drift_fn`	f(x, t) -> dx/dt. The drift function of the SDE.	required
`diffusion_matrix`	B, the diffusion matrix (constant). Shape (d, m).	required
`in_dim`	State dimension.	required
`materialize`	Kept for API consistency — Phi is always dense here because expm requires dense matrices.	`False`

Returns:

Name	Type	Description
`transition_model_fn`		(x, cov, t_old, t) -> (x_pred, Phi, Q)

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

def make_continuous_transition(drift_fn, diffusion_matrix, in_dim, materialize=False):
    """Build an EKF transition model from continuous SDE.

    For an SDE  dx = f(x,t) dt + B dw,  this computes:
      - x_pred = x + f(x, t_old) * (t - t_old)   (Euler step)
      - Phi, Q via matrix_fraction_decomposition

    Since matrix_fraction_decomposition must materialize the 2d x 2d block
    matrix for expm, Phi is always dense regardless of the materialize flag.
    The flag only controls whether the separate Jacobian A is kept dense
    or discarded (it's computed regardless for the decomposition).

    Args:
        drift_fn: f(x, t) -> dx/dt.  The drift function of the SDE.
        diffusion_matrix: B, the diffusion matrix (constant). Shape (d, m).
        in_dim: State dimension.
        materialize: Kept for API consistency — Phi is always dense here
            because expm requires dense matrices.

    Returns:
        transition_model_fn: (x, cov, t_old, t) -> (x_pred, Phi, Q)
    """
    B = diffusion_matrix

    def transition_model_fn(x, cov, t_old, t):
        dt = t - t_old
        x_pred = x + drift_fn(x, t_old) * dt

        # Dense Jacobian required for matrix_fraction_decomposition
        A = jax.jacfwd(lambda x: drift_fn(x, t_old))(x)
        Phi, Q = matrix_fraction_decomposition(t_old, t, A, B)
        return x_pred, Phi, Q

    return transition_model_fn

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

    @staticmethod
    def default_unpack(state, info):
        return (state.mean, state.cov)

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
    sigma_points_L_2L = mu0 - sqrt_cov
    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
    sigma_points_neg = mu0 - sqrt_cov
    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]
    sqrt_val = jnp.sqrt(D + 1)

    # Create a (D+1, D) array filled with -√(D+1)
    base_points = jnp.full((D + 1, D), -sqrt_val)

    # Create row and column indices
    rows = jnp.arange(D + 1)[:, None]  # Shape (D+1,1)
    cols = jnp.arange(D)[None, :]  # Shape (1,D)

    # Set a diagonal pattern: for (i+1, i), set to +√(D+1)
    mask = (rows - 1) == cols
    base_points = jnp.where(mask, sqrt_val, base_points)

    # Ensure the points have zero mean
    base_points = base_points - jnp.mean(base_points, axis=0, keepdims=True)

    # Scale to achieve unit covariance before transforming by cov0
    scale = jnp.sqrt(D / (2 * (D + 1)))
    base_points = base_points * scale

    # Apply the covariance transformation
    A = jnp.linalg.cholesky(cov0)
    sigma_points = mu0[None, :] + base_points @ A

    # 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
        )
        if isinstance(transition_covariance_matrix, Callable):
            Q = transition_covariance_matrix(t_old, t)
        else:
            Q = transition_covariance_matrix

        # Prediction step
        mu1_, cov1_ = unscented_transform(
            predicted_sigma_points, weights_mean, weights_cov, noise_cov=Q
        )

        if is_observed:
            # 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=observation_covariance,
            )

            # 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 = -0.5 * (
                jnp.linalg.slogdet(cov_y)[1] + r.T @ jnp.linalg.solve(cov_y, r)
            )

            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 ¶

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))

build_kernel ¶

build_kernel(transition_model_fns, observation_model_fns, rank, linear_solve=None, min_eig=1e-09, process_noise_rank=None, energy_threshold=None, min_rank=1)

Build a rank-reduced Kalman filter kernel.

Covariance is represented as P ~= U S U^T with U in R^{d x r}, S in R^{r x r}. Predict step is performed in low-rank form by combining propagated state covariance and process covariance through a compressed factorization. Update step uses low-rank algebra in the current subspace.

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

def build_kernel(
    transition_model_fns: Callable,
    observation_model_fns: Callable,
    rank: int,
    linear_solve: Optional[Callable] = None,
    min_eig: float = 1e-9,
    process_noise_rank: Optional[int] = None,
    energy_threshold: Optional[float] = None,
    min_rank: int = 1,
) -> Callable:
    """Build a rank-reduced Kalman filter kernel.

    Covariance is represented as P ~= U S U^T with U in R^{d x r}, S in R^{r x r}.
    Predict step is performed in low-rank form by combining propagated state
    covariance and process covariance through a compressed factorization.
    Update step uses low-rank algebra in the current subspace.
    """

    if process_noise_rank is None:
        process_noise_rank = rank

    if energy_threshold is not None and not (0.0 < energy_threshold <= 1.0):
        raise ValueError("energy_threshold must be in (0, 1].")
    min_rank = max(1, min(min_rank, rank))

    def kernel(
        state: RankReducedKalmanFilterState,
        t: Optional[ArrayLike] = None,
        observed: Optional[ArrayLike] = None,
        rng_key: Optional[jnp.ndarray] = None,
    ) -> Tuple[RankReducedKalmanFilterState, RankReducedKalmanFilterInfo]:
        del rng_key

        mu0 = state.mean
        U0 = state.cov_factor
        S0 = state.cov_core
        t_old = state.t
        is_observed = observed is not None

        transition = transition_model_fns(t_old, t)
        if isinstance(transition, tuple) and len(transition) == 3:
            Phi, Uq, Sq = transition
            Uq, Sq = _parse_process_noise(
                (Uq, Sq),
                process_noise_rank=process_noise_rank,
                min_eig=min_eig,
            )
        else:
            Phi, Q = transition
            Uq, Sq = _parse_process_noise(
                Q,
                process_noise_rank=process_noise_rank,
                min_eig=min_eig,
            )

        # Predict in low-rank form: P_pred = Phi U0 S0 U0^T Phi^T + Uq Sq Uq^T
        Up = Phi @ U0
        U_pred, S_pred = _sum_lowrank_terms(
            Up,
            S0,
            Uq,
            Sq,
            rank=rank,
            min_eig=min_eig,
        )
        pred_vals = jnp.diag(S_pred)
        pred_vals = _apply_energy_threshold(
            pred_vals,
            min_eig=min_eig,
            energy_threshold=energy_threshold,
            min_rank=min_rank,
        )
        S_pred = jnp.diag(pred_vals)
        mu_pred = Phi @ mu0

        if is_observed:
            C, R = observation_model_fns(t)
            y = observed
            y_pred = C @ mu_pred
            residual = y - y_pred

            # Innovation: S_y = C U S U^T C^T + R
            M = C @ U_pred
            innovation = M @ S_pred @ M.T + R

            # K = U S M^T S_y^{-1}
            B = S_pred @ M.T
            solve = _default_solve if linear_solve is None else linear_solve
            B_Sinv = solve(innovation, B)
            K = U_pred @ B_Sinv

            mu = mu_pred + K @ residual

            # Covariance update in reduced coordinates:
            # S_post = S - S M^T S_y^{-1} M S
            S_post = symmetrize_matrix(S_pred - B_Sinv @ (M @ S_pred))

            # Re-orthogonalize/truncate in-subspace if needed
            vals, vecs = _eigh_clip(S_post, min_eig=min_eig)
            k = min(rank, vals.shape[0])
            vals = vals[-k:]
            vecs = vecs[:, -k:]
            vals_desc = vals[::-1]
            vals_desc = _apply_energy_threshold(
                vals_desc,
                min_eig=min_eig,
                energy_threshold=energy_threshold,
                min_rank=min_rank,
            )
            vals = vals_desc[::-1]
            U = U_pred @ vecs
            S = jnp.diag(vals)

            logdet = jnp.linalg.slogdet(jnp.asarray(innovation))[1]
            log_likelihood = -0.5 * (logdet + residual.T @ solve(innovation, residual))

            return RankReducedKalmanFilterState(
                mu, U, S, t
            ), RankReducedKalmanFilterInfo(
                mu_pred,
                U_pred,
                S_pred,
                log_likelihood,
            )

        return RankReducedKalmanFilterState(
            mu_pred, U_pred, S_pred, t
        ), RankReducedKalmanFilterInfo(mu_pred, U_pred, S_pred, jnp.array(0.0))

    return kernel

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: run an existing sampler in a flow's latent space.

A posterior with strong curvature or correlation is hard for HMC not because the kernel is weak but because the geometry is bad -- one step size cannot suit every direction. NeuTra fixes the geometry instead of the kernel: fit a flow T to the target, then sample the pulled-back density

log p~(z) = log p(T(z)) + log|det J_T(z)|

which is close to an isotropic Gaussian whenever the flow is any good, and push the draws back through T.

The key property: this is a change of variables, not an approximation. MCMC on p~ remains asymptotically exact for p no matter how poor the flow is. A bad flow costs efficiency, never correctness -- which is what makes it safe to pair with the mode-seeking reverse-KL fit in :mod:probjax.inference.vi.flow_vi.

Measured on Neal's funnel (D=5, 4000 draws, matched budget): NUTS on the transformed target reached ESS 601 against 168 for NUTS on the target directly.

Nothing here is a new kernel, so every existing kernel, warmup and runner works unchanged:

transform = neutra(logdensity_fn, flow) kernel = nuts(transform.logdensity) # or mala, hmc, mclmc, slice... state = kernel.init(key, jnp.zeros(dim)) result = MCMC(kernel).sample(key, state, 4000, kernel.init_params(state)) draws = transform.forward(result.samples) # back in the target's space

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

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

Source code in probjax/inference/base.py

class MCMCResult(NamedTuple):
    state: Any
    samples: Optional[Any] = None
    info: Optional[Any] = None

probjax.inference.SMCResult ¶

Bases: NamedTuple

Source code in probjax/inference/base.py

class SMCResult(NamedTuple):
    state: Any
    params: Any
    info: Optional[Any] = None
    log_evidence: Optional[Any] = None
    final_info: Optional[Any] = None
    num_steps: Optional[Any] = None
    completed: Optional[Any] = None
    key: Optional[Any] = None

probjax.inference.FilteringResult ¶

Bases: NamedTuple

Source code in probjax/inference/base.py

class FilteringResult(NamedTuple):
    initial_state: Any
    states: Any
    info: Optional[Any] = None

probjax.inference.AdaptationResult ¶

Bases: NamedTuple

Source code in probjax/inference/base.py

class AdaptationResult(NamedTuple):
    state: Any
    params: Any
    trace: Optional[Any] = None
    final_info: Optional[Any] = None

probjax.inference.WarmupResult ¶

Bases: NamedTuple

Source code in probjax/inference/base.py

class WarmupResult(NamedTuple):
    state: Any
    params: Any
    info: Optional[Any] = None

probjax.inference.MarkovKernel `module-attribute` ¶

MarkovKernel = Kernel

probjax.inference.Kernel ¶

Bases: NamedTuple

A pure transition and the functions needed to initialize it.

Source code in probjax/inference/base.py

class Kernel(NamedTuple):
    """A pure transition and the functions needed to initialize it."""

    init: Callable
    step: Callable
    init_params: Callable

    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 ¶

sample_kernel staticmethod ¶

run ¶

sample ¶

adapt ¶

warmup ¶

probjax.inference.SMC ¶

run_adaptive_kernel staticmethod ¶

MCMC kernels¶

probjax.inference.hmc ¶

init_params ¶

build_step ¶

probjax.inference.nuts module-attribute ¶

probjax.inference.dynamic_hmc ¶

probjax.inference.mala ¶

probjax.inference.mclmc ¶

probjax.inference.adjusted_mclmc module-attribute ¶

probjax.inference.mh ¶

probjax.inference.gauss_rwmh module-attribute ¶

probjax.inference.imh ¶

NeuralIMHParams ¶

NeuralIMHState ¶

NeuralIMHWarmupInfo ¶

GaussianIMHParams ¶

wrap_logpdf ¶

build_imh_step ¶

init_imh_params ¶

neural_imh ¶

neural_imh_warmup ¶

init_gaussian_imh_params ¶

proposal_gaussian ¶

proposal_gaussian_logpdf ¶

probjax.inference.neural_imh ¶

probjax.inference.slice ¶

probjax.inference.latent_slice ¶

probjax.inference.elliptical_slice ¶

probjax.inference.arms ¶

probjax.inference.a2rms module-attribute ¶

probjax.inference.ars ¶

probjax.inference.gaussian_imh module-attribute ¶

probjax.inference.adjusted_mclmc_dynamic module-attribute ¶

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.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 ¶

Sequential Monte Carlo¶

probjax.inference.smc ¶

smc ¶

persistent_smc_kernel ¶

adaptive_persistent_smc_kernel ¶

probjax.inference.adaptive_smc module-attribute ¶

probjax.inference.persistent_smc ¶

probjax.inference.adaptive_persistent_smc ¶

probjax.inference.path_smc ¶

probjax.inference.GeometricPath dataclass ¶

probjax.inference.PartialPosteriorsPath dataclass ¶

SMC extensions¶

probjax.inference.waste_free_strategy ¶

probjax.inference.waste_free_smc ¶

probjax.inference.tuned_smc ¶

run_kernel `staticmethod` ¶

sample_kernel `staticmethod` ¶

run_adaptive_kernel `staticmethod` ¶

probjax.inference.nuts `module-attribute` ¶

probjax.inference.adjusted_mclmc `module-attribute` ¶

probjax.inference.gauss_rwmh `module-attribute` ¶

probjax.inference.a2rms `module-attribute` ¶

probjax.inference.gaussian_imh `module-attribute` ¶

probjax.inference.adjusted_mclmc_dynamic `module-attribute` ¶

probjax.inference.adapt `module-attribute` ¶

probjax.inference.adaptive_smc `module-attribute` ¶

probjax.inference.GeometricPath `dataclass` ¶

probjax.inference.PartialPosteriorsPath `dataclass` ¶