Hamiltonian HMC code with PyMC JAX - GPU sampler

[andrej]

@jessegrabowski updating jnp.pi to np.pi as in the two codes above did not affect the error, as it seems
However thank you for your suggestion