# Path: blackjax/types.py



from typing import Callable, NamedTuple, Tuple
from blackjax.types import Array, PRNGKey, PyTree
import jax
import blackjax.mcmc.integrators as integrators
import blackjax.mcmc.metrics as metrics
import blackjax.mcmc.proposal as proposal
import blackjax.mcmc.trajectory as trajectory
"""Public API for the HMC Kernel"""



__all__ = ["HMCState", "HMCInfo", "init", "kernel"]


class HMCState(NamedTuple):
    """State of the HMC algorithm.

    The HMC algorithm takes one position of the chain and returns another
    position. In order to make computations more efficient, we also store
    the current potential energy as well as the current gradient of the
    potential energy.
    """


# Path: blackjax/types.py



# Path: blackjax/mcmc/hmc.py
from typing import Callable, NamedTuple, Tuple
from blackjax.types import Array, PRNGKey, PyTree
import jax
import blackjax.mcmc.integrators as integrators
import blackjax.mcmc.metrics as metrics
import blackjax.mcmc.proposal as proposal
import blackjax.mcmc.trajectory as trajectory

"""Public API for the HMC Kernel"""



__all__ = ["HMCState", "HMCInfo", "init", "kernel"]


class HMCState(NamedTuple):
    """State of the HMC algorithm.

    The HMC algorithm takes one position of the chain and returns another
    position. In order to make computations more efficient, we also store
    the current potential energy as well as the current gradient of the
    potential energy.
    """


    position: PyTree