AeHMC provides MCMC sampling algorithms written in Aesara.
- Sample from an (unnormalized) probability distribution using Hamiltonian Monte Carlo and the No U-Turn Sampler.
import aesara
from aesara import tensor as at
from aesara.tensor.random.utils import RandomStream
from aeppl import joint_logprob
from aehmc import nuts
# A simple normal distribution
Y_rv = at.random.normal(0, 1)
def logprob_fn(y):
logprob = joint_logprob({Y_rv: y})
return logprob
# Build the transition kernel
srng = RandomStream(seed=0)
kernel = nuts.kernel(
srng,
logprob_fn,
inverse_mass_matrix=at.as_tensor(1.),
)
# Compile a function that updates the chain
y_vv = Y_rv.clone()
initial_state = nuts.new_state(y_vv, logprob_fn)
next_step = kernel(*initial_state, 1e-2)
print(next_step[0][1].eval({y_vv: 0}))The latest release of AeHMC can be installed from PyPI using pip:
pip install aehmc
Or via conda-forge:
conda install -c conda-forge aehmc
The current development branch of AeHMC can be installed from GitHub using pip:
pip install git+https://github.com/aesara-devs/aehmc