Repository: pyro-ppl/numpyro

Issue:
Clarify that `sample(..., obs_mask)` is only used for SVI
to avoid confusion as in this [forum thread](https://forum.pyro.ai/t/bayesian-imputation/3443).


Discussion:

Resolving patch:
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -128,7 +128,8 @@ def sample(
         broadcastable with ``fn.batch_shape``. If provided, events with
         mask=True will be conditioned on ``obs`` and remaining events will be
         imputed by sampling. This introduces a latent sample site named ``name
-        + "_unobserved"`` which should be used by guides.
+        + "_unobserved"`` which should be used by guides in SVI. Note that this
+        argument is not intended to be used with MCMC.
     :return: sample from the stochastic `fn`.
     """
     assert isinstance(


Test patch:
diff --git a/test/infer/test_compute_downstream_costs.py b/test/infer/test_compute_downstream_costs.py
--- a/test/infer/test_compute_downstream_costs.py
+++ b/test/infer/test_compute_downstream_costs.py
@@ -360,10 +360,10 @@ def test_compute_downstream_costs_plate_reuse(dim1, dim2):
 
     for k in dc:
         assert guide_trace[k]["log_prob"].shape == dc[k].shape
-        assert_allclose(dc[k], dc_brute[k])
+        assert_allclose(dc[k], dc_brute[k], rtol=1e-6)
 
     expected_c1 = model_trace["c1"]["log_prob"] - guide_trace["c1"]["log_prob"]
     expected_c1 += (model_trace["b1"]["log_prob"] - guide_trace["b1"]["log_prob"]).sum()
     expected_c1 += model_trace["c2"]["log_prob"] - guide_trace["c2"]["log_prob"]
     expected_c1 += model_trace["obs"]["log_prob"]
-    assert_allclose(expected_c1, dc["c1"])
+    assert_allclose(expected_c1, dc["c1"], rtol=1e-6)
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
@@ -7,8 +7,8 @@
 from numpy.testing import assert_allclose
 import pytest
 
+import jax
 from jax import device_get, jit, lax, pmap, random, vmap
-from jax.lib import xla_bridge
 import jax.numpy as jnp
 from jax.scipy.special import logit
 from jax.test_util import check_close
@@ -586,7 +586,7 @@ def model(labels):
 def test_chain_inside_jit(kernel_cls, chain_method):
     # NB: this feature is useful for consensus MC.
     # Caution: compiling time will be slow (~ 90s)
-    if chain_method == "parallel" and xla_bridge.device_count() == 1:
+    if chain_method == "parallel" and jax.device_count() == 1:
         pytest.skip("parallel method requires device_count greater than 1.")
     num_warmup, num_samples = 100, 2000
     # Here are settings which is currently supported.
@@ -716,7 +716,7 @@ def model(data):
     reason="without this mark, we have duplicated tests in Travis",
 )
 def test_functional_map(algo, map_fn):
-    if map_fn is pmap and xla_bridge.device_count() == 1:
+    if map_fn is pmap and jax.device_count() == 1:
         pytest.skip("pmap test requires device_count greater than 1.")
 
     true_mean, true_std = 1.0, 2.0
diff --git a/test/test_compile.py b/test/test_compile.py
--- a/test/test_compile.py
+++ b/test/test_compile.py
@@ -3,8 +3,8 @@
 
 import pytest
 
+import jax
 from jax import lax, random
-from jax.lib import xla_bridge
 
 import numpyro
 import numpyro.distributions as dist
@@ -42,9 +42,7 @@ def test_mcmc_one_chain(deterministic, find_heuristic_step_size):
 
 
 @pytest.mark.parametrize("deterministic", [True, False])
-@pytest.mark.skipif(
-    xla_bridge.device_count() < 2, reason="only one device is available"
-)
+@pytest.mark.skipif(jax.device_count() < 2, reason="only one device is available")
 def test_mcmc_parallel_chain(deterministic):
     GLOBAL["count"] = 0
     mcmc = MCMC(NUTS(model), num_warmup=100, num_samples=100, num_chains=2)