def hmc_move(s_rng, positions, energy_fn, stepsize, n_steps):
 initial_vel = s_rng.normal(size=positions.shape)
 (final_pos, final_vel) = simulate_dynamics(initial_pos=positions, initial_vel=initial_vel, stepsize=stepsize, n_steps=n_steps, energy_fn=energy_fn)
 accept = metropolis_hastings_accept(energy_prev=hamiltonian(positions, initial_vel, energy_fn), energy_next=hamiltonian(final_pos, final_vel, energy_fn), s_rng=s_rng)
 return (accept, final_pos)

def hmc_move ( s_rng , positions , energy_fn , stepsize , n_steps ) : initial_vel = s_rng . normal ( size = positions . shape ) ( final_pos , final_vel ) = simulate_dynamics ( initial_pos = positions , initial_vel = initial_vel , stepsize = stepsize , n_steps = n_steps , energy_fn = energy_fn ) accept = metropolis_hastings_accept ( energy_prev = hamiltonian ( positions , initial_vel , energy_fn ) , energy_next = hamiltonian ( final_pos , final_vel , energy_fn ) , s_rng = s_rng ) return ( accept , final_pos )
this function performs one-step of hybrid monte-carlo sampling .