Fixing issue Samples are outside the support for DiscreteUniform distribution pyro-ppl#1834

Author: Deathn0t

numpyro/infer/hmc_gibbs.py

@@ -434,6 +434,13 @@ def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):
             and site["fn"].has_enumerate_support
             and not site["is_observed"]
         }
+        self._support_enumerates = {
+            name: site["fn"].enumerate_support(False)
+            for name, site in self._prototype_trace.items()
+            if site["type"] == "sample"
+            and site["fn"].has_enumerate_support
+            and not site["is_observed"]
+        }
         self._gibbs_sites = [
             name
             for name, site in self._prototype_trace.items()

numpyro/infer/mixed_hmc.py

@@ -6,6 +6,7 @@
 
 from jax import grad, jacfwd, lax, random
 from jax.flatten_util import ravel_pytree
+import jax
 import jax.numpy as jnp
 
 from numpyro.infer.hmc import momentum_generator
@@ -301,6 +302,11 @@ def body_fn(i, vals):
             adapt_state=adapt_state,
         )
 
+        z_discrete = jax.tree.map(
+            lambda idx, support: support[idx],
+            z_discrete,
+            self._support_enumerates,
+        )
         z = {**z_discrete, **hmc_state.z}
         return MixedHMCState(z, hmc_state, rng_key, accept_prob)