@@ -0,0 +1,2653 @@
+%22%22%22Proposals in the HMC family.%0A%0AProposals take the current state of the chain, transform it to another state%0Athat is returned to the base kernel. Proposals can differ from one another on%0Atwo aspects: the way the step size is chosen and the way the number of%0Aintegration steps is chose.%0A%0AThe standard HMC algorithm integrates the same number of times with the same%0Astep size %5B1%5D_. It is also common to draw at each step the number of%0Aintegration steps from a distribution %5B1,2%5D_ ; empirical HMC %5B2%5D_ for instance%0Alearns this distribution during the adaptation phase. Other algorithms, like%0ANUTS %5B3, 4, 5%5D_, determine the number of integration steps dynamically at runtime.%0A%0AReferences%0A----------%0A.. %5B1%5D: Duane, Simon, et al. %22Hybrid monte carlo.%22 Physics letters B 195.2 (1987): 216-222.%0A.. %5B2%5D: Wu, Changye, Julien Stoehr, and Christian P. Robert. %22Faster%0A        Hamiltonian Monte Carlo by learning leapfrog scale.%22 arXiv preprint%0A        arXiv:1810.04449 (2018).%0A.. %5B3%5D: Hoffman, Matthew D., and Andrew Gelman. %22The No-U-Turn sampler:%0A        adaptively setting path lengths in Hamiltonian Monte Carlo.%22 J. Mach. Learn.%0A        Res. 15.1 (2014): 1593-1623.%0A.. %5B4%5D: Phan, Du, Neeraj Pradhan, and Martin Jankowiak. %22Composable effects for%0A        flexible and accelerated probabilistic programming in NumPyro.%22 arXiv preprint%0A        arXiv:1912.11554 (2019).%0A.. %5B5%5D: Lao, Junpeng, et al. %22tfp. mcmc: Modern Markov Chain Monte Carlo Tools%0A        Built for Modern Hardware.%22 arXiv preprint arXiv:2002.01184 (2020).%0A%0A%22%22%22%0Afrom typing import Callable, NamedTuple, Tuple%0A%0Aimport jax%0Aimport jax.numpy as jnp%0A%0Afrom blackjax.inference.integrators import IntegratorState%0A%0A__all__ = %5B%22hmc%22%5D%0A%0AHMCProposalState = IntegratorState%0A%0A%0Aclass HMCProposalInfo(NamedTuple):%0A    step_size: float%0A    num_integration_steps: int%0A    traversed_states: IntegratorState  # intermediate + last state%0A%0A%0Adef hmc(%0A    integrator: Callable, step_size: float, num_integration_steps: int = 1%0A) -%3E Callable:%0A    %22%22%22Vanilla HMC proposal running the integrator for a fixed number of steps%22%22%22%0A%0A    def propose(%0A        _, initial_state: HMCProposalState%0A    ) -%3E Tuple%5BHMCProposalState, HMCProposalInfo%5D:%0A        %22%22%22Integrate the trajectory  %60num_integration_steps%60 times starting from %60initial_state%60.%22%22%22%0A%0A        def one_step(state, _):%0A            state = integrator(state, step_size)%0A            return state, state%0A%0A        new_state, traversed_states = jax.lax.scan(%0A            one_step, initial_state, jnp.arange(num_integration_steps)%0A        )%0A        info = HMCProposalInfo(step_size, num_integration_steps, traversed_states)%0A%0A        return new_state, info%0A%0A    return propose%0A

add hmc proposal