# Path: blackjax/mcmc/hmc.py
# 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
#     potential_energy: float
#     potential_energy_grad: PyTree
# 
# Path: blackjax/types.py



from typing import Callable, NamedTuple, Tuple
from blackjax.mcmc.hmc import HMCState
from blackjax.types import PRNGKey
import jax
import jax.numpy as jnp
import blackjax.adaptation.optimizers as optimizers
        p_accept = info.acceptance_probability
        gradient = target - p_accept
        return DualAveragingAdaptationState(*da_update(da_state, gradient))

    def final(da_state: DualAveragingAdaptationState) -> float:
        return jnp.exp(da_state.log_step_size_avg)

    return init, update, final


# -------------------------------------------------------------------
#                 REASONABLE FIRST STEP SIZE
# -------------------------------------------------------------------


class ReasonableStepSizeState(NamedTuple):
    """State carried through the search for a reasonable first step size.

    rng_key
        Key used by JAX's random number generator.
    direction: {-1, 1}
        Determines whether the step size should be increased or decreased during the
        previous step search. If direction = 1 it will be increased, otherwise decreased.
    previous_direction
        The previous direction. It is necessary to carry it because the choice of step
        size is made at the end of the search update.
    step_size
        The current step size in the search.
    """


# Path: blackjax/mcmc/hmc.py
# 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
#     potential_energy: float
#     potential_energy_grad: PyTree
# 
# Path: blackjax/types.py



# Path: blackjax/adaptation/step_size.py
from typing import Callable, NamedTuple, Tuple
from blackjax.mcmc.hmc import HMCState
from blackjax.types import PRNGKey
import jax
import jax.numpy as jnp
import blackjax.adaptation.optimizers as optimizers

        p_accept = info.acceptance_probability
        gradient = target - p_accept
        return DualAveragingAdaptationState(*da_update(da_state, gradient))

    def final(da_state: DualAveragingAdaptationState) -> float:
        return jnp.exp(da_state.log_step_size_avg)

    return init, update, final


# -------------------------------------------------------------------
#                 REASONABLE FIRST STEP SIZE
# -------------------------------------------------------------------


class ReasonableStepSizeState(NamedTuple):
    """State carried through the search for a reasonable first step size.

    rng_key
        Key used by JAX's random number generator.
    direction: {-1, 1}
        Determines whether the step size should be increased or decreased during the
        previous step search. If direction = 1 it will be increased, otherwise decreased.
    previous_direction
        The previous direction. It is necessary to carry it because the choice of step
        size is made at the end of the search update.
    step_size
        The current step size in the search.
    """


    rng_key: PRNGKey