Repository: pyro-ppl/numpyro

Issue:
handlers.mask not working when enable_validation(True)
Using `handlers.mask` for sampling doesn't work when `numpyro.enable_validation(True)`. 

Is this a bug or is it expected behaviour?

```python
import numpyro
numpyro.enable_validation(True) # all works fine when this is False
from jax import numpy as jnp, random
from numpyro import sample
from numpyro.handlers import mask
from numpyro.distributions import Bernoulli, Beta
from numpyro.infer.mcmc import MCMC
from numpyro.infer.hmc import NUTS

ptrue = 0.5
Y     = random.bernoulli(random.PRNGKey(0), ptrue, (100,))
Y     = jnp.where(jnp.arange(100) < 20, -1, Y)
Yobs  = Y >= 0

def m(Y, Yobs):
    theta = sample('theta', Beta(1,1))
    with numpyro.plate('obs', Y.shape[0]):
        with mask(mask_array=Yobs):
            sample('obs_Y', Bernoulli(probs=theta), obs=Y)

mcmc = MCMC(NUTS(m), 100, 500)
mcmc.run(random.PRNGKey(1), Y, Yobs)

mcmc.print_summary()
```


Discussion:
I think the issue comes from: JAX has changed `0 x inf = 0` to `0 x inf = nan`.

Resolving patch:
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -311,7 +311,8 @@ def mask(self, mask):
         broadcastable to the distributions
         :attr:`Distribution.batch_shape` .
 
-        :param mask: A boolean or boolean valued array.
+        :param mask: A boolean or boolean valued array (`True` includes
+            a site, `False` excludes a site).
         :type mask: bool or jnp.ndarray
         :return: A masked copy of this distribution.
         :rtype: :class:`MaskedDistribution`
@@ -636,7 +637,7 @@ def log_prob(self, value):
             return jnp.zeros(shape)
         if self._mask is True:
             return self.base_dist.log_prob(value)
-        return self.base_dist.log_prob(value) * self._mask
+        return jnp.where(self._mask, self.base_dist.log_prob(value), 0.)
 
     def enumerate_support(self, expand=True):
         return self.base_dist.enumerate_support(expand=expand)
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -373,8 +373,8 @@ class mask(Messenger):
     """
     This messenger masks out some of the sample statements elementwise.
 
-    :param mask: a DeviceArray with `bool` dtype for masking elementwise masking
-        of sample sites.
+    :param mask: a boolean or a boolean-valued array for masking elementwise log
+        probability of sample sites (`True` includes a site, `False` excludes a site).
     """
     def __init__(self, fn=None, mask=True, mask_array=None):
         if mask_array is not None:


Test patch:
diff --git a/test/test_handlers.py b/test/test_handlers.py
--- a/test/test_handlers.py
+++ b/test/test_handlers.py
@@ -45,6 +45,15 @@ def model(data, mask):
     assert_allclose(log_joint, expected, atol=1e-4)
 
 
+def test_mask_inf():
+    def model():
+        with handlers.mask(mask=jnp.zeros(10, dtype=bool)):
+            numpyro.factor('inf', -jnp.inf)
+
+    log_joint = log_density(model, (), {}, {})[0]
+    assert_allclose(log_joint, 0.)
+
+
 @pytest.mark.parametrize('use_context_manager', [True, False])
 def test_scale(use_context_manager):
     def model(data):