An HMC-within-Gibbs sampler for Numpyro
This package adds a new HMC-within-Gibbs sampler to Numpyro. Unlike the HMCGibbs sampler currently available, this sampler is for situations where you do not have an analytic form for one of your conditioned distributions. Instead, it uses an HMC/NUTS sampler to estimate draws from each of the conditioned distributions.
To use MultiHMCGibbs you need to create a list of HMC or NUTS kernels that wrap the same model, but each can have its own keywords such as target_accept_prob or max_tree_depth. The other argument is a list of lists containing the free parameters for each of the inner kernels.
Internally the sampler will:
- Loop over the kernels in the list
- Conditioned it on the non-free parameters
- Re-calculate the likelihood and gradients at the new conditioned point
- Step the kernel forward
- Move on to the next kernel