# Path: blackjax/types.py



from typing import NamedTuple, Tuple
from typing_extensions import Protocol
from .types import PRNGKey, PyTree
    def __call__(self, position: Position) -> State:
        """Initialize the algorithm's state.

        Parameters
        ----------
        position
           A chain position.

        Returns
        -------
        The kernel state that corresponds to the position.

        """


class UpdateFn(Protocol):
    """A transition kernel used as the `update` of a `SamplingAlgorithms`.

    Kernels are pure functions and are idempotent. They necessarily take a
    random state `rng_key` and the current kernel state (which contains the
    current position) as parameters, return a new state and some information
    about the transtion.

    Update functions is a simplified yet universal interface with every sampling
    algorithm. In essence, what all these algorithms do is take a rng state, a
    chain state (possibly a batch of data) and return a new state and some
    information about the transition.

    """


# Path: blackjax/types.py



# Path: blackjax/base.py
from typing import NamedTuple, Tuple
from typing_extensions import Protocol
from .types import PRNGKey, PyTree

    def __call__(self, position: Position) -> State:
        """Initialize the algorithm's state.

        Parameters
        ----------
        position
           A chain position.

        Returns
        -------
        The kernel state that corresponds to the position.

        """


class UpdateFn(Protocol):
    """A transition kernel used as the `update` of a `SamplingAlgorithms`.

    Kernels are pure functions and are idempotent. They necessarily take a
    random state `rng_key` and the current kernel state (which contains the
    current position) as parameters, return a new state and some information
    about the transtion.

    Update functions is a simplified yet universal interface with every sampling
    algorithm. In essence, what all these algorithms do is take a rng state, a
    chain state (possibly a batch of data) and return a new state and some
    information about the transition.

    """


    def __call__(self, rng_key: PRNGKey, state: State) -> Tuple[State, Info]: