Repository: pyro-ppl/numpyro

Issue:
UnexpectedTracerError with DiscreteHMCGibbs, a latent discrete variable with more than 1 dim, and num_chains>1
The combination of 
1) A latent discrete variable with event_shape > 1x
2) num_chains>1

Throws the following error
#UnexpectedTracerError: Encountered an unexpected tracer. Perhaps this tracer escaped through global state from a previously traced function.
#The functions being transformed should not save traced values to global state. Detail: tracer created on line /home/skynet3/numpyro/numpyro/infer/hmc_gibbs.py:356 (<dictcomp>).

Minimally reproducible example
```
import os
import numpyro
import jax
from jax import random
import jax.numpy as jnp
import numpyro.distributions as dist
from numpyro.infer import DiscreteHMCGibbs, MCMC, NUTS
numpyro.set_platform("gpu") 
numpyro.set_host_device_count(2) 
print(jax.__version__) #0.2.8
print(numpyro.__version__) #0.5.0
print(jax.config.FLAGS.jax_backend_target) #local
print(jax.lib.xla_bridge.get_backend().platform) #gpu

#Fails
#UnexpectedTracerError: Encountered an unexpected tracer. Perhaps this tracer escaped through global state from a previously traced function.
#The functions being transformed should not save traced values to global state. Detail: tracer created on line /home/skynet3/numpyro/numpyro/infer/hmc_gibbs.py:356 (<dictcomp>).
def model():
  C = numpyro.sample("C", dist.Bernoulli(probs=0.5 * jnp.ones(5 )  ))
kernel = DiscreteHMCGibbs(NUTS(model), modified=True)
mcmc   = MCMC(kernel, num_warmup=100, num_samples=100, num_chains=2)
mcmc.warmup(random.PRNGKey(0))
mcmc.run(random.PRNGKey(1))
```

The error does not appear if the discrete variable is size 1 or if the number of chains=1. Only the combination of both throws the error

```
#Works
def model():
  C = numpyro.sample("C", dist.Bernoulli(probs=0.5  ))
kernel = DiscreteHMCGibbs(NUTS(model), modified=True)
mcmc   = MCMC(kernel, num_warmup=100, num_samples=100, num_chains=2)
mcmc.warmup(random.PRNGKey(0))
mcmc.run(random.PRNGKey(1))

#Works
def model():
  C = numpyro.sample("C", dist.Bernoulli(probs=0.5 * jnp.ones(5 )  ))
kernel = DiscreteHMCGibbs(NUTS(model), modified=True)
mcmc   = MCMC(kernel, num_warmup=100, num_samples=100, num_chains=1)
mcmc.warmup(random.PRNGKey(0))
mcmc.run(random.PRNGKey(1))
```



Raise errors when there are no gibbs sites in DiscreteHMCGibbs and HMCECS
Currently, if there is no discrete latent site, DiscreteHMCGibbs will throw `OverflowError: cannot convert float infinity to integer` (see #889), which is misleading. We should raise a better error message when there is no gibbs sites in those classes.


Discussion:
Thanks, @rexdouglass, for the detailed explanation! I think we should change [jnp.broadcast_to](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/infer/hmc_gibbs.py#L356) to `(import numpy as np) np.broadcast_to` because `jnp.foo` will create a tracer, rather than concrete arrays when running multiple chains. Do you want to submit a fix together with the above regression test? :) 
I can confirm that changing that single line and adding the import for numpy as np resolves the error on both my small test example above as well as that larger covid nowcast I'm working on.

```
name: np.broadcast_to(site["fn"].enumerate_support(False).shape[0], jnp.shape(site["value"])) #changed jnp to np per discussion here https://github.com/pyro-ppl/numpyro/issues/907
```

I did not change the jnp.shape to np, I don't think it matters unless some later jax bug is introduced about mixing and matching.



Resolving patch:
diff --git a/numpyro/infer/hmc_gibbs.py b/numpyro/infer/hmc_gibbs.py
--- a/numpyro/infer/hmc_gibbs.py
+++ b/numpyro/infer/hmc_gibbs.py
@@ -6,6 +6,8 @@
 from functools import partial
 import warnings
 
+import numpy as np
+
 from jax import device_put, grad, hessian, jacfwd, jacobian, lax, ops, random, value_and_grad
 import jax.numpy as jnp
 from jax.scipy.special import expit
@@ -294,13 +296,11 @@ class DiscreteHMCGibbs(HMCGibbs):
         corresponding :func:`~numpyro.primitives.sample` statement.
 
     :param inner_kernel: One of :class:`~numpyro.infer.hmc.HMC` or :class:`~numpyro.infer.hmc.NUTS`.
-    :param list discrete_sites: a list of site names for the discrete latent variables
-        that are covered by the Gibbs sampler.
     :param bool random_walk: If False, Gibbs sampling will be used to draw a sample from the
         conditional `p(gibbs_site | remaining sites)`. Otherwise, a sample will be drawn uniformly
-        from the domain of `gibbs_site`.
+        from the domain of `gibbs_site`. Defaults to False.
     :param bool modified: whether to use a modified proposal, as suggested in reference [1], which
-        always proposes a new state for the current Gibbs site.
+        always proposes a new state for the current Gibbs site. Defaults to False.
         The modified scheme appears in the literature under the name "modified Gibbs sampler" or
         "Metropolised Gibbs sampler".
 
@@ -328,7 +328,7 @@ class DiscreteHMCGibbs(HMCGibbs):
         >>> kernel = DiscreteHMCGibbs(NUTS(model), modified=True)
         >>> mcmc = MCMC(kernel, 1000, 100000, progress_bar=False)
         >>> mcmc.run(random.PRNGKey(0), probs, locs)
-        >>> mcmc.print_summary()
+        >>> mcmc.print_summary()  # doctest: +SKIP
         >>> samples = mcmc.get_samples()["x"]
         >>> assert abs(jnp.mean(samples) - 1.3) < 0.1
         >>> assert abs(jnp.var(samples) - 4.36) < 0.5
@@ -356,7 +356,7 @@ def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):
         self._prototype_trace = trace(seed(self.model, key_u)).get_trace(*model_args, **model_kwargs)
 
         self._support_sizes = {
-            name: jnp.broadcast_to(site["fn"].enumerate_support(False).shape[0], jnp.shape(site["value"]))
+            name: np.broadcast_to(site["fn"].enumerate_support(False).shape[0], jnp.shape(site["value"]))
             for name, site in self._prototype_trace.items()
             if site["type"] == "sample" and site["fn"].has_enumerate_support and not site["is_observed"]
         }
@@ -365,6 +365,7 @@ def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):
                              and site["fn"].has_enumerate_support
                              and not site["is_observed"]
                              and site["infer"].get("enumerate", "") != "parallel"]
+        assert self._gibbs_sites, "Cannot detect any discrete latent variables in the model."
         return super().init(rng_key, num_warmup, init_params, model_args, model_kwargs)
 
     def sample(self, state, model_args, model_kwargs):
@@ -533,6 +534,7 @@ def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):
             if site["type"] == "plate" and site["args"][0] > site["args"][1]  # i.e. size > subsample_size
         }
         self._gibbs_sites = list(self._subsample_plate_sizes.keys())
+        assert self._gibbs_sites, "Cannot detect any subsample statements in the model."
         if self._proxy is not None:
             proxy_fn, gibbs_init, self._gibbs_update = self._proxy(self._prototype_trace,
                                                                    self._subsample_plate_sizes,
diff --git a/numpyro/infer/hmc_util.py b/numpyro/infer/hmc_util.py
--- a/numpyro/infer/hmc_util.py
+++ b/numpyro/infer/hmc_util.py
@@ -415,7 +415,9 @@ def _update_at_window_end(z_info, rng_key_ss, state):
 
         if adapt_step_size:
             step_size = find_reasonable_step_size(step_size, inverse_mass_matrix, z_info, rng_key_ss)
-            ss_state = ss_init(jnp.log(10 * step_size))
+            # NB: when step_size is large, say 1e38, jnp.log(10 * step_size) will be inf
+            # and jnp.log(10) + jnp.log(step_size) will be finite
+            ss_state = ss_init(jnp.log(10) + jnp.log(step_size))
 
         return HMCAdaptState(step_size, inverse_mass_matrix, mass_matrix_sqrt,
                              ss_state, mm_state, window_idx, rng_key)
@@ -475,6 +477,8 @@ def _is_turning(inverse_mass_matrix, r_left, r_right, r_sum):
     elif inverse_mass_matrix.ndim == 1:
         v_left = jnp.multiply(inverse_mass_matrix, r_left)
         v_right = jnp.multiply(inverse_mass_matrix, r_right)
+    else:
+        raise ValueError("inverse_mass_matrix should have 1 or 2 dimensions.")
 
     # This implements dynamic termination criterion (ref [2], section A.4.2).
     r_sum = r_sum - (r_left + r_right) / 2
@@ -737,6 +741,8 @@ def euclidean_kinetic_energy(inverse_mass_matrix, r):
         v = jnp.matmul(inverse_mass_matrix, r)
     elif inverse_mass_matrix.ndim == 1:
         v = jnp.multiply(inverse_mass_matrix, r)
+    else:
+        raise ValueError("inverse_mass_matrix should have 1 or 2 dimensions.")
 
     return 0.5 * jnp.dot(v, r)
 
@@ -748,6 +754,8 @@ def _euclidean_kinetic_energy_grad(inverse_mass_matrix, r):
         v = jnp.matmul(inverse_mass_matrix, r)
     elif inverse_mass_matrix.ndim == 1:
         v = jnp.multiply(inverse_mass_matrix, r)
+    else:
+        raise ValueError("inverse_mass_matrix should have 1 or 2 dimensions.")
 
     return unravel_fn(v)
 


Test patch:
diff --git a/test/test_hmc_gibbs.py b/test/test_hmc_gibbs.py
--- a/test/test_hmc_gibbs.py
+++ b/test/test_hmc_gibbs.py
@@ -150,13 +150,15 @@ def model():
     assert_allclose(x1_std, np.sqrt(np.diagonal(cov11)), rtol=0.1)
 
 
-def test_discrete_gibbs_multiple_sites():
+@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 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, progress_bar=False)
+    mcmc = MCMC(kernel, 1000, 10000, num_chains=num_chains, progress_bar=False)
     mcmc.run(random.PRNGKey(0))
     samples = mcmc.get_samples()
     assert_allclose(jnp.mean(samples["x"], 0), 0.7 * jnp.ones(3), atol=0.01)
@@ -202,7 +204,7 @@ def model(probs, locs):
     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.1)
+    assert_allclose(jnp.var(samples["x"]), 4.36, atol=0.2)
     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/test_hmc_util.py b/test/test_hmc_util.py
--- a/test/test_hmc_util.py
+++ b/test/test_hmc_util.py
@@ -322,7 +322,7 @@ def find_reasonable_step_size(step_size, m_inv, z, rng_key):
     assert window_idx == 3
     # during the last window, because target_accept_prob=0.8,
     # log_step_size will be equal to the constant prox_center=log(10*last_step_size)
-    assert_allclose(step_size, last_step_size * 10)
+    assert_allclose(step_size, last_step_size * 10, atol=1e-6)
     # Verifies that inverse_mass_matrix does not change during the last window
     # despite z_flat changes w.r.t time t,
     assert_allclose(final_inverse_mass_matrix, inverse_mass_matrix)