[FR] Mixed HMC for models with discrete latent variables

[fehiepsi]

Currently, NumPyro uses enumeration to marginalize out discrete latent variables. However, when there are many discrete latent variables, using enumeration will be slow because the dimensional size is blown up exponentially w.r.t. the number of discrete latent variables. It seems to me that @StannisZhou's Mixed HMC [1] is a good inference algorithm to use in that case (e.g. the model in #769, which is pretty slow under enumeration).

References:

• Mixed Hamiltonian Monte Carlo for Mixed Discrete and Continuous Variables
• Discontinuous Hamiltonian Monte Carlo for discrete parameters and discontinuous likelihoods
• JAX implementation for Mixed HMC: https://github.com/StannisZhou/mixed_hmc