# Path: blackjax/types.py



from typing import Callable, Tuple
from blackjax.types import PRNGKey
import jax
import jax.numpy as jnp
import blackjax.smc.base as base
import blackjax.smc.ess as ess
import blackjax.smc.solver as solver
import blackjax.smc.tempered as tempered
    A callable that takes a rng_key and a TemperedSMCState that contains the current state
    of the chain and that returns a new state of the chain along with
    information about the transition.
    """

    def compute_delta(state: tempered.TemperedSMCState) -> float:
        lmbda = state.lmbda
        max_delta = 1 - lmbda
        delta = ess.ess_solver(
            jax.vmap(loglikelihood_fn),
            state.particles,
            target_ess,
            max_delta,
            root_solver,
            use_log_ess,
        )
        delta = jnp.clip(delta, 0.0, max_delta)

        return delta

    kernel = tempered.kernel(
        logprior_fn,
        loglikelihood_fn,
        mcmc_kernel_factory,
        make_mcmc_state,
        resampling_fn,
        mcmc_iter,
    )

    def one_step(

# Path: blackjax/types.py



# Path: blackjax/smc/adaptive_tempered.py
from typing import Callable, Tuple
from blackjax.types import PRNGKey
import jax
import jax.numpy as jnp
import blackjax.smc.base as base
import blackjax.smc.ess as ess
import blackjax.smc.solver as solver
import blackjax.smc.tempered as tempered

    A callable that takes a rng_key and a TemperedSMCState that contains the current state
    of the chain and that returns a new state of the chain along with
    information about the transition.
    """

    def compute_delta(state: tempered.TemperedSMCState) -> float:
        lmbda = state.lmbda
        max_delta = 1 - lmbda
        delta = ess.ess_solver(
            jax.vmap(loglikelihood_fn),
            state.particles,
            target_ess,
            max_delta,
            root_solver,
            use_log_ess,
        )
        delta = jnp.clip(delta, 0.0, max_delta)

        return delta

    kernel = tempered.kernel(
        logprior_fn,
        loglikelihood_fn,
        mcmc_kernel_factory,
        make_mcmc_state,
        resampling_fn,
        mcmc_iter,
    )

    def one_step(

        rng_key: PRNGKey, state: tempered.TemperedSMCState