Repository: pyro-ppl/numpyro

Issue:
[FR] Mixed HMC for models with discrete latent variables
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](https://arxiv.org/abs/1909.04852)
+ [Discontinuous Hamiltonian Monte Carlo for discrete parameters and discontinuous likelihoods](https://arxiv.org/abs/1705.08510)
+ JAX implementation for Mixed HMC: https://github.com/StannisZhou/mixed_hmc


Discussion:

Resolving patch:
diff --git a/numpyro/infer/__init__.py b/numpyro/infer/__init__.py
--- a/numpyro/infer/__init__.py
+++ b/numpyro/infer/__init__.py
@@ -13,6 +13,7 @@
 )
 from numpyro.infer.barker import BarkerMH
 from numpyro.infer.mcmc import MCMC
+from numpyro.infer.mixed_hmc import MixedHMC
 from numpyro.infer.sa import SA
 from numpyro.infer.svi import SVI
 from numpyro.infer.util import Predictive, log_likelihood
@@ -35,6 +36,7 @@
     'HMCECS',
     'HMCGibbs',
     'MCMC',
+    'MixedHMC',
     'NUTS',
     'Predictive',
     'RenyiELBO',
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -22,7 +22,8 @@
 from numpyro.infer.util import ParamInfo, init_to_uniform, initialize_model
 from numpyro.util import cond, fori_loop, identity
 
-HMCState = namedtuple('HMCState', ['i', 'z', 'z_grad', 'potential_energy', 'energy', 'num_steps', 'accept_prob',
+HMCState = namedtuple('HMCState', ['i', 'z', 'z_grad', 'potential_energy', 'energy',
+                                   'r', 'trajectory_length', 'num_steps', 'accept_prob',
                                    'mean_accept_prob', 'diverging', 'adapt_state', 'rng_key'])
 """
 A :func:`~collections.namedtuple` consisting of the following fields:
@@ -33,6 +34,10 @@
  - **z_grad** - Gradient of potential energy w.r.t. latent sample sites.
  - **potential_energy** - Potential energy computed at the given value of ``z``.
  - **energy** - Sum of potential energy and kinetic energy of the current state.
+ - **r** - The current momentum variable. If this is None, a new momentum variable
+   will be drawn at the beginning of each sampling step.
+ - **trajectory_length** - The amount of time to run HMC dynamics in each sampling step.
+   This field is not used in NUTS.
  - **num_steps** - Number of steps in the Hamiltonian trajectory (for diagnostics).
    In NUTS sampler, the tree depth of a trajectory can be computed from this field
    with `tree_depth = np.log2(num_steps).astype(int) + 1`.
@@ -56,7 +61,7 @@
 
 
 def _get_num_steps(step_size, trajectory_length):
-    num_steps = jnp.clip(trajectory_length / step_size, a_min=1)
+    num_steps = jnp.ceil(trajectory_length / step_size)
     # NB: casting to jnp.int64 does not take effect (returns jnp.int32 instead)
     # if jax_enable_x64 is False
     return num_steps.astype(canonicalize_dtype(jnp.int64))
@@ -146,7 +151,6 @@ def hmc(potential_fn=None, potential_fn_gen=None, kinetic_fn=None, algo='NUTS'):
     if kinetic_fn is None:
         kinetic_fn = euclidean_kinetic_energy
     vv_update = None
-    trajectory_len = None
     max_treedepth = None
     wa_update = None
     wa_steps = None
@@ -205,10 +209,9 @@ def init_kernel(init_params,
 
         """
         step_size = lax.convert_element_type(step_size, canonicalize_dtype(jnp.float64))
-        nonlocal wa_update, trajectory_len, max_treedepth, vv_update, wa_steps, forward_mode_ad
+        nonlocal wa_update, max_treedepth, vv_update, wa_steps, forward_mode_ad
         forward_mode_ad = forward_mode_differentiation
         wa_steps = num_warmup
-        trajectory_len = trajectory_length
         max_treedepth = max_tree_depth
         if isinstance(init_params, ParamInfo):
             z, pe, z_grad = init_params
@@ -246,17 +249,24 @@ def init_kernel(init_params,
         vv_state = vv_init(z, r, potential_energy=pe, z_grad=z_grad)
         energy = kinetic_fn(wa_state.inverse_mass_matrix, vv_state.r)
         hmc_state = HMCState(jnp.array(0), vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy,
+                             None, trajectory_length,
                              jnp.array(0), jnp.array(0.), jnp.array(0.), jnp.array(False), wa_state, rng_key_hmc)
         return device_put(hmc_state)
 
     def _hmc_next(step_size, inverse_mass_matrix, vv_state,
-                  model_args, model_kwargs, rng_key):
+                  model_args, model_kwargs, rng_key, trajectory_length):
         if potential_fn_gen:
             nonlocal vv_update, forward_mode_ad
             pe_fn = potential_fn_gen(*model_args, **model_kwargs)
             _, vv_update = velocity_verlet(pe_fn, kinetic_fn, forward_mode_ad)
 
-        num_steps = _get_num_steps(step_size, trajectory_len)
+        # no need to spend too many steps if the state z has 0 size (i.e. z is empty)
+        if inverse_mass_matrix.shape[0] == 0:
+            num_steps = 1
+        else:
+            num_steps = _get_num_steps(step_size, trajectory_length)
+        # makes sure trajectory length is constant, rather than step_size * num_steps
+        step_size = trajectory_length / num_steps
         vv_state_new = fori_loop(0, num_steps,
                                  lambda i, val: vv_update(step_size, inverse_mass_matrix, val),
                                  vv_state)
@@ -273,7 +283,7 @@ def _hmc_next(step_size, inverse_mass_matrix, vv_state,
         return vv_state, energy, num_steps, accept_prob, diverging
 
     def _nuts_next(step_size, inverse_mass_matrix, vv_state,
-                   model_args, model_kwargs, rng_key):
+                   model_args, model_kwargs, rng_key, trajectory_length):
         if potential_fn_gen:
             nonlocal vv_update, forward_mode_ad
             pe_fn = potential_fn_gen(*model_args, **model_kwargs)
@@ -307,14 +317,16 @@ def sample_kernel(hmc_state, model_args=(), model_kwargs=None):
         """
         model_kwargs = {} if model_kwargs is None else model_kwargs
         rng_key, rng_key_momentum, rng_key_transition = random.split(hmc_state.rng_key, 3)
-        r = momentum_generator(hmc_state.z, hmc_state.adapt_state.mass_matrix_sqrt, rng_key_momentum)
+        r = momentum_generator(hmc_state.z, hmc_state.adapt_state.mass_matrix_sqrt, rng_key_momentum) \
+            if hmc_state.r is None else hmc_state.r
         vv_state = IntegratorState(hmc_state.z, r, hmc_state.potential_energy, hmc_state.z_grad)
         vv_state, energy, num_steps, accept_prob, diverging = _next(hmc_state.adapt_state.step_size,
                                                                     hmc_state.adapt_state.inverse_mass_matrix,
                                                                     vv_state,
                                                                     model_args,
                                                                     model_kwargs,
-                                                                    rng_key_transition)
+                                                                    rng_key_transition,
+                                                                    hmc_state.trajectory_length)
         # not update adapt_state after warmup phase
         adapt_state = cond(hmc_state.i < wa_steps,
                            (hmc_state.i, accept_prob, vv_state, hmc_state.adapt_state),
@@ -326,7 +338,9 @@ def sample_kernel(hmc_state, model_args=(), model_kwargs=None):
         n = jnp.where(hmc_state.i < wa_steps, itr, itr - wa_steps)
         mean_accept_prob = hmc_state.mean_accept_prob + (accept_prob - hmc_state.mean_accept_prob) / n
 
-        return HMCState(itr, vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy, num_steps,
+        r = vv_state.r if hmc_state.r is not None else None
+        return HMCState(itr, vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy,
+                        r, hmc_state.trajectory_length, num_steps,
                         accept_prob, mean_accept_prob, diverging, adapt_state, rng_key)
 
     # Make `init_kernel` and `sample_kernel` visible from the global scope once
diff --git a/numpyro/infer/mixed_hmc.py b/numpyro/infer/mixed_hmc.py
new file mode 100644
--- /dev/null
+++ b/numpyro/infer/mixed_hmc.py
@@ -0,0 +1,222 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import namedtuple
+from functools import partial
+
+from jax import jacfwd, grad, lax, ops, random
+import jax.numpy as jnp
+
+from numpyro.infer.hmc import momentum_generator
+from numpyro.infer.hmc_gibbs import DiscreteHMCGibbs
+from numpyro.infer.hmc_util import euclidean_kinetic_energy, warmup_adapter
+from numpyro.util import cond, fori_loop, identity, ravel_pytree
+
+
+MixedHMCState = namedtuple("MixedHMCState", "z, hmc_state, rng_key, accept_prob")
+
+
+class MixedHMC(DiscreteHMCGibbs):
+    """
+    Implementation of Mixed Hamiltonian Monte Carlo (reference [1]).
+
+    .. note:: The number of discrete sites to update at each MCMC iteration
+        (`n_D` in reference [1]) is fixed at value 1.
+
+    **References**
+
+    1. *Mixed Hamiltonian Monte Carlo for Mixed Discrete and Continuous Variables*,
+       Guangyao Zhou (2020)
+    2. *Peskun's theorem and a modified discrete-state Gibbs sampler*,
+       Liu, J. S. (1996)
+
+    :param inner_kernel: A :class:`~numpyro.infer.hmc.HMC` kernel.
+    :param int num_discrete_updates: Number of times to update discrete variables.
+        Defaults to the number of discrete latent variables.
+    :param bool random_walk: If False, Gibbs sampling will be used to draw a sample from the
+        conditional `p(gibbs_site | remaining sites)`, where `gibbs_site` is one of the
+        discrete sample sites in the model. Otherwise, a sample will be drawn uniformly
+        from the domain of `gibbs_site`. Defaults to False.
+    :param bool modified: whether to use a modified proposal, as suggested in reference [2], which
+        always proposes a new state for the current Gibbs site (i.e. discrete site). Defaults to False.
+        The modified scheme appears in the literature under the name "modified Gibbs sampler" or
+        "Metropolised Gibbs sampler".
+
+    **Example**
+
+    .. doctest::
+
+        >>> from jax import random
+        >>> import jax.numpy as jnp
+        >>> import numpyro
+        >>> import numpyro.distributions as dist
+        >>> from numpyro.infer import HMC, MCMC, MixedHMC
+        ...
+        >>> def model(probs, locs):
+        ...     c = numpyro.sample("c", dist.Categorical(probs))
+        ...     numpyro.sample("x", dist.Normal(locs[c], 0.5))
+        ...
+        >>> probs = jnp.array([0.15, 0.3, 0.3, 0.25])
+        >>> locs = jnp.array([-2, 0, 2, 4])
+        >>> kernel = MixedHMC(HMC(model, trajectory_length=1.2), num_discrete_updates=20)
+        >>> mcmc = MCMC(kernel, 1000, 100000, progress_bar=False)
+        >>> mcmc.run(random.PRNGKey(0), probs, locs)
+        >>> mcmc.print_summary()  # doctest: +SKIP
+        >>> samples = mcmc.get_samples()
+        >>> assert "x" in samples and "c" in samples
+        >>> assert abs(jnp.mean(samples["x"]) - 1.3) < 0.1
+        >>> assert abs(jnp.var(samples["x"]) - 4.36) < 0.5
+    """
+
+    def __init__(self, inner_kernel, *, num_discrete_updates=None, random_walk=False, modified=False):
+        super().__init__(inner_kernel, random_walk=random_walk, modified=modified)
+        if inner_kernel._algo == "NUTS":
+            raise ValueError("The algorithm only works with HMC and and does not support NUTS.")
+        self._num_discrete_updates = num_discrete_updates
+
+    def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):
+        rng_key, rng_r = random.split(rng_key)
+        state = super().init(rng_key, num_warmup, init_params, model_args, model_kwargs)
+        self._support_sizes_flat, _ = ravel_pytree({k: self._support_sizes[k] for k in self._gibbs_sites})
+        if self._num_discrete_updates is None:
+            self._num_discrete_updates = self._support_sizes_flat.shape[0]
+        self._num_warmup = num_warmup
+
+        # NB: the warmup adaptation can not be performed in sub-trajectories (i.e. the hmc trajectory
+        # between two discrete updates), so we will do it here, at the end of each MixedHMC step.
+        _, self._wa_update = warmup_adapter(num_warmup,
+                                            adapt_step_size=self.inner_kernel._adapt_step_size,
+                                            adapt_mass_matrix=self.inner_kernel._adapt_mass_matrix,
+                                            dense_mass=self.inner_kernel._dense_mass,
+                                            target_accept_prob=self.inner_kernel._target_accept_prob,
+                                            find_reasonable_step_size=None)
+
+        # In HMC, when `hmc_state.r` is not None, we will skip drawing a random momemtum at the
+        # beginning of an HMC step. The reason is we need to maintain `r` between each sub-trajectories.
+        r = momentum_generator(state.hmc_state.z, state.hmc_state.adapt_state.mass_matrix_sqrt, rng_r)
+        return MixedHMCState(state.z, state.hmc_state._replace(r=r), state.rng_key, jnp.array(0.))
+
+    def sample(self, state, model_args, model_kwargs):
+        model_kwargs = {} if model_kwargs is None else model_kwargs
+        num_discretes = self._support_sizes_flat.shape[0]
+
+        def potential_fn(z_gibbs, z_hmc):
+            return self.inner_kernel._potential_fn_gen(
+                *model_args, _gibbs_sites=z_gibbs, **model_kwargs)(z_hmc)
+
+        def update_discrete(idx, rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum):
+            # Algo 1, line 19: get a new discrete proposal
+            rng_key, z_discrete_new, pe_new, log_accept_ratio = self._discrete_proposal_fn(
+                rng_key, z_discrete, hmc_state.potential_energy,
+                partial(potential_fn, z_hmc=hmc_state.z), idx, self._support_sizes_flat[idx])
+            # Algo 1, line 20: depending on reject or refract, we will update
+            # the discrete variable and its corresponding kinetic energy. In case of
+            # refract, we will need to update the potential energy and its grad w.r.t. hmc_state.z
+            ke_discrete_i_new = ke_discrete[idx] + log_accept_ratio
+            grad_ = jacfwd if self.inner_kernel._forward_mode_differentiation else grad
+            z_discrete, pe, ke_discrete_i, z_grad = lax.cond(
+                ke_discrete_i_new > 0,
+                (z_discrete_new, pe_new, ke_discrete_i_new),
+                lambda vals: vals + (grad_(partial(potential_fn, vals[0]))(hmc_state.z),),
+                (z_discrete, hmc_state.potential_energy, ke_discrete[idx], hmc_state.z_grad),
+                identity)
+
+            delta_pe_sum = delta_pe_sum + pe - hmc_state.potential_energy
+            ke_discrete = ops.index_update(ke_discrete, idx, ke_discrete_i)
+            hmc_state = hmc_state._replace(potential_energy=pe, z_grad=z_grad)
+            return rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum
+
+        def update_continuous(hmc_state, z_discrete):
+            model_kwargs_ = model_kwargs.copy()
+            model_kwargs_["_gibbs_sites"] = z_discrete
+            hmc_state_new = self.inner_kernel.sample(hmc_state, model_args, model_kwargs_)
+
+            # each time a sub-trajectory is performed, we need to reset i and adapt_state
+            # (we will only update them at the end of HMCGibbs step)
+            # For `num_steps`, we will record its cumulative sum for diagnostics
+            hmc_state = hmc_state_new._replace(i=hmc_state.i,
+                                               adapt_state=hmc_state.adapt_state,
+                                               num_steps=hmc_state.num_steps + hmc_state_new.num_steps)
+            return hmc_state
+
+        def body_fn(i, vals):
+            rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum, arrival_times = vals
+            idx = jnp.argmin(arrival_times)
+            # NB: length of each sub-trajectory is scaled from the current min(arrival_times)
+            # (see the note at total_time below)
+            trajectory_length = arrival_times[idx] * time_unit
+            arrival_times = arrival_times - arrival_times[idx]
+            arrival_times = ops.index_update(arrival_times, idx, 1.)
+
+            # this is a trick, so that in a sub-trajectory of HMC, we always accept the new proposal
+            pe = jnp.inf
+            hmc_state = hmc_state._replace(trajectory_length=trajectory_length, potential_energy=pe)
+            # Algo 1, line 7: perform a sub-trajectory
+            hmc_state = update_continuous(hmc_state, z_discrete)
+            # Algo 1, line 8: perform a discrete update
+            rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum = update_discrete(
+                idx, rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum)
+            return rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum, arrival_times
+
+        z_discrete = {k: v for k, v in state.z.items() if k not in state.hmc_state.z}
+        rng_key, rng_ke, rng_time, rng_r, rng_accept = random.split(state.rng_key, 5)
+        # Algo 1, line 2: sample discrete kinetic energy
+        ke_discrete = random.exponential(rng_ke, (num_discretes,))
+        # Algo 1, line 4 and 5: sample the initial amount of time that each discrete site visits
+        # the point 0/1. The logic in GetStepSizesNSteps(...) is more complicated but does
+        # the same job: the sub-trajectory length eta_t * M_t is the lag between two arrival time.
+        arrival_times = random.uniform(rng_time, (num_discretes,))
+        # compute the amount of time to make `num_discrete_updates` discrete updates
+        total_time = (self._num_discrete_updates - 1) // num_discretes \
+            + jnp.sort(arrival_times)[(self._num_discrete_updates - 1) % num_discretes]
+        # NB: total_time can be different from the HMC trajectory length, so we need to scale
+        # the time unit so that total_time * time_unit = hmc_trajectory_length
+        time_unit = state.hmc_state.trajectory_length / total_time
+
+        # Algo 1, line 2: sample hmc momentum
+        r = momentum_generator(state.hmc_state.r, state.hmc_state.adapt_state.mass_matrix_sqrt, rng_r)
+        hmc_state = state.hmc_state._replace(r=r, num_steps=0)
+        hmc_ke = euclidean_kinetic_energy(hmc_state.adapt_state.inverse_mass_matrix, r)
+        # Algo 1, line 10: compute the initial energy
+        energy_old = hmc_ke + hmc_state.potential_energy
+
+        # Algo 1, line 3: set initial values
+        delta_pe_sum = 0.
+        init_val = (rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum, arrival_times)
+        # Algo 1, line 6-9: perform the update loop
+        rng_key, hmc_state_new, z_discrete_new, _, delta_pe_sum, _ = fori_loop(
+            0, self._num_discrete_updates, body_fn, init_val)
+        # Algo 1, line 10: compute the proposal energy
+        hmc_ke = euclidean_kinetic_energy(hmc_state.adapt_state.inverse_mass_matrix, hmc_state_new.r)
+        energy_new = hmc_ke + hmc_state_new.potential_energy
+        # Algo 1, line 11: perform MH correction
+        delta_energy = energy_new - energy_old - delta_pe_sum
+        delta_energy = jnp.where(jnp.isnan(delta_energy), jnp.inf, delta_energy)
+        accept_prob = jnp.clip(jnp.exp(-delta_energy), a_max=1.0)
+
+        # record the correct new num_steps
+        hmc_state = hmc_state._replace(num_steps=hmc_state_new.num_steps)
+        # reset the trajectory length
+        hmc_state_new = hmc_state_new._replace(trajectory_length=hmc_state.trajectory_length)
+        hmc_state, z_discrete = cond(random.bernoulli(rng_key, accept_prob),
+                                     (hmc_state_new, z_discrete_new), identity,
+                                     (hmc_state, z_discrete), identity)
+
+        # perform hmc adapting (similar to the implementation in hmc)
+        adapt_state = cond(hmc_state.i < self._num_warmup,
+                           (hmc_state.i, accept_prob, (hmc_state.z,), hmc_state.adapt_state),
+                           lambda args: self._wa_update(*args),
+                           hmc_state.adapt_state,
+                           identity)
+
+        itr = hmc_state.i + 1
+        n = jnp.where(hmc_state.i < self._num_warmup, itr, itr - self._num_warmup)
+        mean_accept_prob_prev = state.hmc_state.mean_accept_prob
+        mean_accept_prob = mean_accept_prob_prev + (accept_prob - mean_accept_prob_prev) / n
+        hmc_state = hmc_state._replace(i=itr,
+                                       accept_prob=accept_prob,
+                                       mean_accept_prob=mean_accept_prob,
+                                       adapt_state=adapt_state)
+
+        z = {**z_discrete, **hmc_state.z}
+        return MixedHMCState(z, hmc_state, rng_key, accept_prob)


Test patch:
diff --git a/test/infer/test_hmc_gibbs.py b/test/infer/test_hmc_gibbs.py
--- a/test/infer/test_hmc_gibbs.py
+++ b/test/infer/test_hmc_gibbs.py
@@ -13,7 +13,7 @@
 
 import numpyro
 import numpyro.distributions as dist
-from numpyro.infer import HMC, HMCECS, MCMC, NUTS, DiscreteHMCGibbs, HMCGibbs
+from numpyro.infer import HMC, HMCECS, MCMC, NUTS, DiscreteHMCGibbs, HMCGibbs, MixedHMC
 from numpyro.infer.util import log_density
 
 
@@ -149,61 +149,78 @@ def model():
     assert_allclose(x1_std, np.sqrt(np.diagonal(cov11)), rtol=0.1)
 
 
+@pytest.mark.parametrize("kernel, inner_kernel, kwargs", [
+    (MixedHMC, HMC, {"num_discrete_updates": 5}),
+    (DiscreteHMCGibbs, NUTS, {})
+])
 @pytest.mark.parametrize("num_chains", [1, 2])
 @pytest.mark.filterwarnings("ignore:There are not enough devices:UserWarning")
-def test_discrete_gibbs_multiple_sites_chain(num_chains):
+def test_discrete_gibbs_multiple_sites_chain(kernel, inner_kernel, kwargs, num_chains):
     def model():
         numpyro.sample("x", dist.Bernoulli(0.7).expand([3]))
         numpyro.sample("y", dist.Binomial(10, 0.3))
 
-    kernel = DiscreteHMCGibbs(NUTS(model))
-    mcmc = MCMC(kernel, 1000, 10000, num_chains=num_chains, progress_bar=False)
+    sampler = kernel(inner_kernel(model), **kwargs)
+    mcmc = MCMC(sampler, 1000, 10000, num_chains=num_chains, progress_bar=False)
     mcmc.run(random.PRNGKey(0))
+    mcmc.print_summary()
     samples = mcmc.get_samples()
     assert_allclose(jnp.mean(samples["x"], 0), 0.7 * jnp.ones(3), atol=0.01)
     assert_allclose(jnp.mean(samples["y"], 0), 0.3 * 10, atol=0.1)
 
 
-def test_discrete_gibbs_enum():
+@pytest.mark.parametrize("kernel, inner_kernel, kwargs", [
+    (MixedHMC, HMC, {"num_discrete_updates": 6}),
+    (DiscreteHMCGibbs, NUTS, {})
+])
+def test_discrete_gibbs_enum(kernel, inner_kernel, kwargs):
     def model():
         numpyro.sample("x", dist.Bernoulli(0.7), infer={"enumerate": "parallel"})
         y = numpyro.sample("y", dist.Binomial(10, 0.3))
         numpyro.deterministic("y2", y ** 2)
 
-    kernel = DiscreteHMCGibbs(NUTS(model))
-    mcmc = MCMC(kernel, 1000, 10000, progress_bar=False)
+    sampler = kernel(inner_kernel(model), **kwargs)
+    mcmc = MCMC(sampler, 1000, 10000, progress_bar=False)
     mcmc.run(random.PRNGKey(0))
     samples = mcmc.get_samples()
     assert_allclose(jnp.mean(samples["y"], 0), 0.3 * 10, atol=0.1)
 
 
 @pytest.mark.parametrize("random_walk", [False, True])
-@pytest.mark.parametrize("modified", [False, True])
-def test_discrete_gibbs_bernoulli(random_walk, modified):
+@pytest.mark.parametrize("kernel, inner_kernel, kwargs", [
+    (MixedHMC, HMC, {"num_discrete_updates": 6}),
+    (DiscreteHMCGibbs, NUTS, {"modified": True}),
+    (DiscreteHMCGibbs, NUTS, {"modified": False}),
+])
+def test_discrete_gibbs_bernoulli(random_walk, kernel, inner_kernel, kwargs):
     def model():
         numpyro.sample("c", dist.Bernoulli(0.8))
 
-    kernel = DiscreteHMCGibbs(NUTS(model), random_walk=random_walk, modified=modified)
-    mcmc = MCMC(kernel, 1000, 200000, progress_bar=False)
+    sampler = kernel(inner_kernel(model), random_walk=random_walk, **kwargs)
+    mcmc = MCMC(sampler, 1000, 10000, progress_bar=False)
     mcmc.run(random.PRNGKey(0))
     samples = mcmc.get_samples()["c"]
     assert_allclose(jnp.mean(samples), 0.8, atol=0.05)
 
 
 @pytest.mark.parametrize("modified", [False, True])
-def test_discrete_gibbs_gmm_1d(modified):
+@pytest.mark.parametrize("kernel, inner_kernel, kwargs", [
+    (MixedHMC, HMC, {"num_discrete_updates": 20}),
+    (DiscreteHMCGibbs, NUTS, {})
+])
+def test_discrete_gibbs_gmm_1d(modified, kernel, inner_kernel, kwargs):
     def model(probs, locs):
         c = numpyro.sample("c", dist.Categorical(probs))
         numpyro.sample("x", dist.Normal(locs[c], 0.5))
 
     probs = jnp.array([0.15, 0.3, 0.3, 0.25])
     locs = jnp.array([-2, 0, 2, 4])
-    kernel = DiscreteHMCGibbs(NUTS(model), modified=modified)
-    mcmc = MCMC(kernel, 1000, 200000, progress_bar=False)
+    sampler = kernel(inner_kernel(model, trajectory_length=1.2), modified=modified, **kwargs)
+    mcmc = MCMC(sampler, 1000, 200000, progress_bar=False)
     mcmc.run(random.PRNGKey(0), probs, locs)
     samples = mcmc.get_samples()
     assert_allclose(jnp.mean(samples["x"]), 1.3, atol=0.1)
-    assert_allclose(jnp.var(samples["x"]), 4.36, atol=0.2)
+    assert_allclose(jnp.var(samples["x"]), 4.36, atol=0.4)
     assert_allclose(jnp.mean(samples["c"]), 1.65, atol=0.1)
     assert_allclose(jnp.var(samples["c"]), 1.03, atol=0.1)
 
diff --git a/test/infer/test_mcmc.py b/test/infer/test_mcmc.py
--- a/test/infer/test_mcmc.py
+++ b/test/infer/test_mcmc.py
@@ -228,7 +228,7 @@ def model():
         numpyro.sample("x", dist.MultivariateNormal(jnp.zeros(2), covariance_matrix=true_cov))
 
     if kernel_cls is HMC or kernel_cls is NUTS:
-        kernel = kernel_cls(model, trajectory_length=1., dense_mass=True)
+        kernel = kernel_cls(model, trajectory_length=2., dense_mass=True)
     elif kernel_cls is BarkerMH:
         kernel = BarkerMH(model, dense_mass=True)