# Path: blackjax/types.py



from typing import Callable, Tuple
from jax.flatten_util import ravel_pytree
from blackjax.types import Array, PRNGKey, PyTree
import jax
import jax.numpy as jnp
import jax.scipy as jscipy

An important particular case (and the most used in practice) of metric for the
position space in the Euclidean metric. It is defined by a definite positive
matrix :math:`M` with fixed value so that the kinetic energy of the hamiltonian
dynamic is independent of the position and only depends on the momentum
:math:`p` [1]_.

For a Newtonian hamiltonian dynamic the kinetic energy is given by:

.. math::
    K(p) = \frac{1}{2} p^T M^{-1} p

We can also generate a relativistic dynamic [2]_.

References
----------
.. [1]: Betancourt, Michael, et al. "The geometric foundations of hamiltonian
        monte carlo." Bernoulli 23.4A (2017): 2257-2298.
.. [2]: Lu, Xiaoyu, et al. "Relativistic monte carlo." Artificial Intelligence and Statistics. PMLR, 2017.

"""



__all__ = ["gaussian_euclidean"]

EuclideanKineticEnergy = Callable[[PyTree], float]


def gaussian_euclidean(

# Path: blackjax/types.py



# Path: blackjax/mcmc/metrics.py
from typing import Callable, Tuple
from jax.flatten_util import ravel_pytree
from blackjax.types import Array, PRNGKey, PyTree
import jax
import jax.numpy as jnp
import jax.scipy as jscipy


An important particular case (and the most used in practice) of metric for the
position space in the Euclidean metric. It is defined by a definite positive
matrix :math:`M` with fixed value so that the kinetic energy of the hamiltonian
dynamic is independent of the position and only depends on the momentum
:math:`p` [1]_.

For a Newtonian hamiltonian dynamic the kinetic energy is given by:

.. math::
    K(p) = \frac{1}{2} p^T M^{-1} p

We can also generate a relativistic dynamic [2]_.

References
----------
.. [1]: Betancourt, Michael, et al. "The geometric foundations of hamiltonian
        monte carlo." Bernoulli 23.4A (2017): 2257-2298.
.. [2]: Lu, Xiaoyu, et al. "Relativistic monte carlo." Artificial Intelligence and Statistics. PMLR, 2017.

"""



__all__ = ["gaussian_euclidean"]

EuclideanKineticEnergy = Callable[[PyTree], float]


def gaussian_euclidean(

    inverse_mass_matrix: Array,