# Path: blackjax/mcmc/integrators.py
# class IntegratorState(NamedTuple):
#     """State of the trajectory integration.
# 
#     We keep the gradient of the potential energy to speedup computations.
#     """
# 
#     position: PyTree
#     momentum: PyTree
#     potential_energy: float
#     potential_energy_grad: PyTree


from typing import Callable, NamedTuple, Tuple
from blackjax.mcmc.integrators import IntegratorState
import jax
import jax.numpy as jnp
import numpy as np




class Proposal(NamedTuple):
    """Proposal for the next chain step.
    state:
        The trajectory state corresponding to this proposal.
    energy:
        The potential energy corresponding to the state.
    weight:
        Weight of the proposal. It is equal to the logarithm of the sum of the canonical
        densities of each state :math:`e^{-H(z)}` along the trajectory.
    sum_log_p_accept:
        cumulated Metropolis-Hastings acceptance probabilty across entire trajectory.
    """


# Path: blackjax/mcmc/integrators.py
# class IntegratorState(NamedTuple):
#     """State of the trajectory integration.
# 
#     We keep the gradient of the potential energy to speedup computations.
#     """
# 
#     position: PyTree
#     momentum: PyTree
#     potential_energy: float
#     potential_energy_grad: PyTree


# Path: blackjax/mcmc/proposal.py
from typing import Callable, NamedTuple, Tuple
from blackjax.mcmc.integrators import IntegratorState
import jax
import jax.numpy as jnp
import numpy as np





class Proposal(NamedTuple):
    """Proposal for the next chain step.
    state:
        The trajectory state corresponding to this proposal.
    energy:
        The potential energy corresponding to the state.
    weight:
        Weight of the proposal. It is equal to the logarithm of the sum of the canonical
        densities of each state :math:`e^{-H(z)}` along the trajectory.
    sum_log_p_accept:
        cumulated Metropolis-Hastings acceptance probabilty across entire trajectory.
    """


    state: IntegratorState