Repository: pyro-ppl/numpyro

Issue:
Possible error in the validation of a Categorical distribution
I am getting  an error when I try to run the following code. The code just sample from a categorical distribution using the defined probabilities. 

```python
import numpyro
import numpyro.distributions as dist
import jax.numpy as jnp
numpyro.enable_validation(True)


def model():
    probs = jnp.array([0.5, 0.5, 0.])
    c = numpyro.sample('c', dist.Categorical(probs=probs))
    return c

with numpyro.handlers.seed(rng_seed=54):
    print(model())
```

```
ValueError                                Traceback (most recent call last)
<ipython-input-1-fc7fe60e083b> in <module>
     10 
     11 with numpyro.handlers.seed(rng_seed=54):
---> 12     print(model())

<ipython-input-1-fc7fe60e083b> in model()
      6 def model():
      7     probs = jnp.array([0.5, 0.5, 0.])
----> 8     c = numpyro.sample('c', dist.Categorical(probs=probs))
      9     return c
     10 

~/miniconda3/envs/numpyro_test/lib/python3.8/site-packages/numpyro/distributions/discrete.py in Categorical(probs, logits, validate_args)
    348 def Categorical(probs=None, logits=None, validate_args=None):
    349     if probs is not None:
--> 350         return CategoricalProbs(probs, validate_args=validate_args)
    351     elif logits is not None:
    352         return CategoricalLogits(logits, validate_args=validate_args)

~/miniconda3/envs/numpyro_test/lib/python3.8/site-packages/numpyro/distributions/discrete.py in __init__(self, probs, validate_args)
    265             raise ValueError("`probs` parameter must be at least one-dimensional.")
    266         self.probs = probs
--> 267         super(CategoricalProbs, self).__init__(batch_shape=jnp.shape(self.probs)[:-1],
    268                                                validate_args=validate_args)
    269 

~/miniconda3/envs/numpyro_test/lib/python3.8/site-packages/numpyro/distributions/distribution.py in __init__(self, batch_shape, event_shape, validate_args)
    142                 if not_jax_tracer(is_valid):
    143                     if not is_valid:
--> 144                         raise ValueError("The parameter {} has invalid values".format(param))
    145         super(Distribution, self).__init__()
    146 

ValueError: The parameter probs has invalid values
```

I think the problem is caused by the validation because If I restart my kernel and comment the line ```numpyro.enable_validation(True)``` the code will run without problem. It will print 0 in my case.


If I write a similar code in Pyro with the validation enabled, I do not get an error.

```python
import torch
import pyro
import pyro.distributions as dist
pyro.enable_validation(True)
pyro.set_rng_seed(54)

def model():
    probs = torch.tensor([0.5, 0.5, 0.])
    c = pyro.sample('c', dist.Categorical(probs=probs))
    return c

print(model())
```

I am using Python 3.8.5, Pyro 1.4.0 and NumPyro 0.3.0 with Ubuntu. Happy to help with what I can.


Discussion:
Thanks, @RaulPL! Could you help us relax the check for interval constraint at [this line](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/distributions/constraints.py#L149)?
I think the line to change is this [one](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/distributions/constraints.py#L195). Go from this: 

```python
return jnp.all(x > 0, axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)
```
to this:
```python
return jnp.all(x >= 0, axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)
```

I think this change will fix it. I can make the pull request if that's okay with you.
You are right, sorry for making confusion. Thanks so much for diagnosing the issue and the fix!

Resolving patch:
diff --git a/numpyro/distributions/constraints.py b/numpyro/distributions/constraints.py
--- a/numpyro/distributions/constraints.py
+++ b/numpyro/distributions/constraints.py
@@ -192,7 +192,7 @@ def __call__(self, x):
 class _Simplex(Constraint):
     def __call__(self, x):
         x_sum = jnp.sum(x, axis=-1)
-        return jnp.all(x > 0, axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)
+        return jnp.all(x >= 0, axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)
 
 
 # TODO: Make types consistent


Test patch: