Inference¶

@partial(
    jax.jit,
    static_argnames=("kernel", "num_steps", "adaptor", "collect", "thin", "progress"),
)
def run(
    key: RngKey,
    kernel: Kernel,
    state: Any,
    params: Any,
    num_steps: int,
    *,
    args: tuple[Any, ...] | None = None,
    adaptor: Adaptor | None = None,
    collect: tuple[str, ...] | None = None,
    thin: int = 1,
    progress: bool | tuple[str, ...] = False,
) -> RunResult:
    """Run ``num_steps`` transitions and keep every ``thin``-th.

    Args:
        key: Random key.
        kernel: The kernel to run.
        state: Initial kernel state.
        params: Kernel parameters.
        num_steps: Number of transitions; must be a multiple of ``thin``.
        args: Per-step extra arguments (e.g. minibatches), a tuple whose
            leaves have a leading axis of length ``num_steps``.
        adaptor: Adapts ``params`` after every transition. The returned
            params are the adaptor's finalized ones.
        collect: State or info fields to record at every kept step;
            ``"state"`` and ``"info"`` record them whole. The
            default records the position of MCMC-style states and nothing
            for population states such as SMC's.
        thin: Keep every ``thin``-th transition.
        progress: Show a progress bar. ``True`` shows the log density and
            acceptance rate; a tuple of names shows those state/info fields.

    Returns:
        A :class:`RunResult` whose ``trace`` maps each name in ``collect`` to
        an array with ``num_steps // thin`` rows.
    """
    if num_steps % thin:
        raise ValueError(f"num_steps={num_steps} must be a multiple of thin={thin}.")
    num_kept = num_steps // thin
    keys = jax.random.split(key, (num_kept, thin))
    if args is None:
        args = ()
    args = jax.tree.map(lambda a: a.reshape((num_kept, thin) + a.shape[1:]), args)
    adaptor_state = None if adaptor is None else adaptor.init(state, params)
    if collect is None:
        collect = ("position",) if hasattr(state, "position") else ()

    names = DEFAULT_PROGRESS_STATS if progress is True else progress or ()
    if names:
        # Show only the stats this kernel has, found from the output shapes.
        first_args = jax.tree.map(lambda a: a[0, 0], args)
        state_shape, info_shape = jax.eval_shape(
            kernel.step, keys[0, 0], state, params, *first_args
        )
        names = tuple(
            name
            for name in names
            if hasattr(state_shape, name) or hasattr(info_shape, name)
        )

    def transition(carry, xs):
        state, params, adaptor_state = carry
        step_key, step_args = xs
        state, params, adaptor_state, info = _step(
            step_key,
            kernel,
            state,
            params,
            *step_args,
            adaptor=adaptor,
            adaptor_state=adaptor_state,
        )
        return (state, params, adaptor_state), info

    def keep_one(carry, xs):
        carry, info = jax.lax.scan(transition, tuple(carry), xs)
        state, last_info = carry[0], jax.tree.map(lambda v: v[-1], info)
        record = {name: _field(name, state, last_info) for name in collect}
        stats = tuple(
            jnp.mean(jnp.asarray(_field(name, state, last_info), jnp.float32))
            for name in names
        )
        return carry, (record, stats)

    carry = (state, params, adaptor_state)
    if names:
        carry, (trace, _) = progress_scan(
            keep_one,
            carry,
            (keys, args),
            num_kept,
            names=names,
            stats_fn=lambda output: output[1],
        )
    else:
        carry, (trace, _) = jax.lax.scan(keep_one, carry, (keys, args))
    state, params, adaptor_state = carry
    if adaptor is not None:
        params, _ = adaptor.finalize(adaptor_state, params)
    return RunResult(state, params, trace if collect else None)

def warmup(
    key: RngKey,
    kernel: Kernel,
    state: Any,
    params: Any,
    num_steps: int,
    policy: Adaptor | Warmup | None = None,
) -> RunResult:
    """Fit the kernel parameters over ``num_steps`` transitions.

    ``policy`` is an :class:`Adaptor` (run with :func:`run`, then finalized) or
    a :class:`Warmup` procedure. ``None`` uses the kernel's
    ``default_adaptation``.
    """
    if policy is None:
        if kernel.default_adaptation is None:
            raise ValueError(
                "This kernel has no default adaptation; pass policy= explicitly."
            )
        policy = kernel.default_adaptation()
    if isinstance(policy, Warmup):
        return policy.run(key, kernel, state, params, num_steps)
    if isinstance(policy, Adaptor):
        return run(key, kernel, state, params, num_steps, adaptor=policy, collect=())
    raise TypeError(
        f"policy must be an Adaptor or a Warmup, got {type(policy).__name__}."
    )

@partial(jax.jit, static_argnames=("kernel", "adaptor"))
def step(
    key: RngKey,
    kernel: Kernel,
    state: Any,
    params: Any,
    *args: Any,
    adaptor: Adaptor | None = None,
    adaptor_state: Any = None,
) -> tuple[Any, Any, Any, Any]:
    """Take one transition, then apply one adaptor update if given.

    Returns ``(state, params, adaptor_state, info)``. Use it to write your own
    outer loop, for example one that alternates sampling with fitting.
    """
    return _step(
        key, kernel, state, params, *args, adaptor=adaptor, adaptor_state=adaptor_state
    )

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

    Args:
        inner_kernel_cls: A probjax kernel class (e.g. ``hmc``, ``mala``,
            ``gauss_rwmh``) or a configured one (e.g.
            ``hmc.configure(num_integration_steps=10)``) that exposes
            ``build_step``, ``init`` and ``init_params``.
        stochastic_logdensity_fn: A callable
            ``(position, rng) -> float`` that returns an unbiased
            estimate of the log-density (or log-marginal-likelihood)
            given auxiliary randomness from *rng*.
        num_samples: Number of independent keys to evaluate and average
            for variance reduction (default 1).
        **inner_kernel_kwargs: Options for the inner kernel. They are routed
            to whichever of ``init``, ``init_params`` and ``build_step``
            accepts them.

    Returns:
        A :class:`Kernel` 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")

    configured = as_configured(inner_kernel_cls, **inner_kernel_kwargs)

    # ------------------------------------------------------------------
    # init: evaluate the stochastic logdensity once to populate state
    # ------------------------------------------------------------------
    def init(rng: RngKey, position: Any):
        key_init, key_logdensity = jax.random.split(rng)
        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 configured.init(
            position, logdensity_fn=fixed_logdensity, rng=key_init
        )

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

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

    # Inherit the inner kernel's recommended warmup, unless it needs the
    # target: a stochastic log-density has no single one to bind.
    default_adaptation = configured.default_adaptation_fn
    if default_adaptation is not None and (
        "logdensity_fn" in inspect.signature(default_adaptation).parameters
    ):
        default_adaptation = None
    default_adaptation = bind_default_adaptation(
        default_adaptation, **inner_kernel_kwargs
    )
    return pseudo_marginal_type(init, step, init_params, default_adaptation)

def bound_adaptor(
    field: str = "log_bound", margin: float = 0.0, axis_name: str | None = None
) -> Adaptor:
    """Raise ``params.<field>`` to the largest ``log p - log q`` seen.

    The bound only grows, so draws taken while it is still growing are
    slightly biased towards where it was too low; run this as warmup and keep
    the bound fixed for the draws you report. With ``axis_name`` the chains
    under ``jax.vmap`` share the maximum over all of them.
    """

    def init(state: Any, params: Any) -> tuple[()]:
        return ()

    def update(state: Any, info: RejectionInfo, adaptor_state: Any, params: Any):
        seen = info.max_log_ratio + margin
        if axis_name is not None:
            seen = jax.lax.pmax(seen, axis_name)
        bound = jnp.maximum(get_param(params, field), seen)
        return adaptor_state, replace_params(params, **{field: bound}), None

    def finalize(adaptor_state: Any, params: Any) -> tuple[Any, None]:
        return params, None

    return Adaptor(init, update, finalize)

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

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

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

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

def sgnht(
    grad_estimator: Callable,
    alpha: float = 0.01,
    beta: float = 0.0,
    temperature: float = 1.0,
) -> Kernel:
    """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 the rng to initialise momentum::

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

    def sgnht_state_init(rng: RngKey, position: Any):
        return _sgnht_init(position, rng=rng, alpha=alpha)

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

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

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

    @partial(jax.jit, static_argnames=("kernel", "num_steps"))
    def run(rng, kernel: Kernel, state: Any, params: Any, num_steps: int) -> RunResult:
        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(rng, 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 RunResult(state, params)

    return Warmup(run)

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
    # BlackJAX types the adaptation procedure's ``run`` with two arguments,
    # but the Pathfinder procedure also takes ``num_steps``.
    procedure: Any = 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=("_kernel", "num_steps"))
    def run(rng, _kernel, state, params, num_steps):
        result, info = procedure.run(rng, state.position, num_steps)
        params = replace_params(params, **result.parameters)
        return RunResult(result.state, params, info if collect else None)

    return Warmup(run)