# Path: blackjax/types.py



from typing import Callable, NamedTuple, Tuple
from blackjax.types import Array, PRNGKey, PyTree
import jax
import jax.numpy as jnp
import numpy as np
import blackjax.mcmc.hmc as hmc
import blackjax.mcmc.integrators as integrators
import blackjax.mcmc.metrics as metrics
import blackjax.mcmc.proposal as proposal
import blackjax.mcmc.termination as termination
import blackjax.mcmc.trajectory as trajectory

class NUTSInfo(NamedTuple):
    """Additional information on the NUTS transition.

    This additional information can be used for debugging or computing
    diagnostics.

    momentum:
        The momentum that was sampled and used to integrate the trajectory.
    is_divergent
        Whether the difference in energy between the original and the new state
        exceeded the divergence threshold.
    is_turning
        Whether the sampling returned because the trajectory started turning
        back on itself.
    energy:
        Energy of the transition.
    trajectory_leftmost_state
        The leftmost state of the full trajectory.
    trajectory_rightmost_state
        The rightmost state of the full trajectory.
    num_trajectory_expansions
        Number of subtrajectory samples that were taken.
    integration_steps
        Number of integration steps that were taken. This is also the number of
        states in the full trajectory.
    acceptance_probability
        average acceptance probabilty across entire trajectory
    """


# Path: blackjax/types.py



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


class NUTSInfo(NamedTuple):
    """Additional information on the NUTS transition.

    This additional information can be used for debugging or computing
    diagnostics.

    momentum:
        The momentum that was sampled and used to integrate the trajectory.
    is_divergent
        Whether the difference in energy between the original and the new state
        exceeded the divergence threshold.
    is_turning
        Whether the sampling returned because the trajectory started turning
        back on itself.
    energy:
        Energy of the transition.
    trajectory_leftmost_state
        The leftmost state of the full trajectory.
    trajectory_rightmost_state
        The rightmost state of the full trajectory.
    num_trajectory_expansions
        Number of subtrajectory samples that were taken.
    integration_steps
        Number of integration steps that were taken. This is also the number of
        states in the full trajectory.
    acceptance_probability
        average acceptance probabilty across entire trajectory
    """


    momentum: PyTree