We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `numpyro/distributions/constraints.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 # The implementation follows the design in PyTorch: torch.distributions.constraints.py
5 #
6 # Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
7 # Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
8 # Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
9 # Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
10 # Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
11 # Copyright (c) 2011-2013 NYU                      (Clement Farabet)
12 # Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
13 # Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
14 # Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
15 #
16 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
17 # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18 # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
19 # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
20 # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
21 # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
22 # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
23 # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
24 # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
25 # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
26 # POSSIBILITY OF SUCH DAMAGE.
27 
28 
29 __all__ = [
30     'boolean',
31     'corr_cholesky',
32     'corr_matrix',
33     'dependent',
34     'greater_than',
35     'integer_interval',
36     'integer_greater_than',
37     'interval',
38     'is_dependent',
39     'less_than',
40     'lower_cholesky',
41     'multinomial',
42     'nonnegative_integer',
43     'positive',
44     'positive_definite',
45     'positive_integer',
46     'real',
47     'real_vector',
48     'simplex',
49     'unit_interval',
50     'Constraint',
51 ]
52 
53 import jax.numpy as jnp
54 
55 
56 class Constraint(object):
57     """
58     Abstract base class for constraints.
59 
60     A constraint object represents a region over which a variable is valid,
61     e.g. within which a variable can be optimized.
62     """
63 
64     def __call__(self, x):
65         raise NotImplementedError
66 
67     def check(self, value):
68         """
69         Returns a byte tensor of `sample_shape + batch_shape` indicating
70         whether each event in value satisfies this constraint.
71         """
72         return self(value)
73 
74 
75 class _Boolean(Constraint):
76     def __call__(self, x):
77         return (x == 0) | (x == 1)
78 
79 
80 class _CorrCholesky(Constraint):
81     def __call__(self, x):
82         tril = jnp.tril(x)
83         lower_triangular = jnp.all(jnp.reshape(tril == x, x.shape[:-2] + (-1,)), axis=-1)
84         positive_diagonal = jnp.all(jnp.diagonal(x, axis1=-2, axis2=-1) > 0, axis=-1)
85         x_norm = jnp.linalg.norm(x, axis=-1)
86         unit_norm_row = jnp.all((x_norm <= 1) & (x_norm > 1 - 1e-6), axis=-1)
87         return lower_triangular & positive_diagonal & unit_norm_row
88 
89 
90 class _CorrMatrix(Constraint):
91     def __call__(self, x):
92         # check for symmetric
93         symmetric = jnp.all(jnp.all(x == jnp.swapaxes(x, -2, -1), axis=-1), axis=-1)
94         # check for the smallest eigenvalue is positive
95         positive = jnp.linalg.eigh(x)[0][..., 0] > 0
96         # check for diagonal equal to 1
97         unit_variance = jnp.all(jnp.abs(jnp.diagonal(x, axis1=-2, axis2=-1) - 1) < 1e-6, axis=-1)
98         return symmetric & positive & unit_variance
99 
100 
101 class _Dependent(Constraint):
102     def __call__(self, x):
103         raise ValueError('Cannot determine validity of dependent constraint')
104 
105 
106 def is_dependent(constraint):
107     return isinstance(constraint, _Dependent)
108 
109 
110 class _GreaterThan(Constraint):
111     def __init__(self, lower_bound):
112         self.lower_bound = lower_bound
113 
114     def __call__(self, x):
115         return x > self.lower_bound
116 
117 
118 class _LessThan(Constraint):
119     def __init__(self, upper_bound):
120         self.upper_bound = upper_bound
121 
122     def __call__(self, x):
123         return x < self.upper_bound
124 
125 
126 class _IntegerInterval(Constraint):
127     def __init__(self, lower_bound, upper_bound):
128         self.lower_bound = lower_bound
129         self.upper_bound = upper_bound
130 
131     def __call__(self, x):
132         return (x >= self.lower_bound) & (x <= self.upper_bound) & (x == jnp.floor(x))
133 
134 
135 class _IntegerGreaterThan(Constraint):
136     def __init__(self, lower_bound):
137         self.lower_bound = lower_bound
138 
139     def __call__(self, x):
140         return (x % 1 == 0) & (x >= self.lower_bound)
141 
142 
143 class _Interval(Constraint):
144     def __init__(self, lower_bound, upper_bound):
145         self.lower_bound = lower_bound
146         self.upper_bound = upper_bound
147 
148     def __call__(self, x):
149         return (x > self.lower_bound) & (x < self.upper_bound)
150 
151 
152 class _LowerCholesky(Constraint):
153     def __call__(self, x):
154         tril = jnp.tril(x)
155         lower_triangular = jnp.all(jnp.reshape(tril == x, x.shape[:-2] + (-1,)), axis=-1)
156         positive_diagonal = jnp.all(jnp.diagonal(x, axis1=-2, axis2=-1) > 0, axis=-1)
157         return lower_triangular & positive_diagonal
158 
159 
160 class _Multinomial(Constraint):
161     def __init__(self, upper_bound):
162         self.upper_bound = upper_bound
163 
164     def __call__(self, x):
165         return jnp.all(x >= 0, axis=-1) & (jnp.sum(x, -1) == self.upper_bound)
166 
167 
168 class _OrderedVector(Constraint):
169     def __call__(self, x):
170         return jnp.all(x[..., 1:] > x[..., :-1], axis=-1)
171 
172 
173 class _PositiveDefinite(Constraint):
174     def __call__(self, x):
175         # check for symmetric
176         symmetric = jnp.all(jnp.all(x == jnp.swapaxes(x, -2, -1), axis=-1), axis=-1)
177         # check for the smallest eigenvalue is positive
178         positive = jnp.linalg.eigh(x)[0][..., 0] > 0
179         return symmetric & positive
180 
181 
182 class _Real(Constraint):
183     def __call__(self, x):
184         return jnp.isfinite(x)
185 
186 
187 class _RealVector(Constraint):
188     def __call__(self, x):
189         return jnp.all(jnp.isfinite(x), axis=-1)
190 
191 
192 class _Simplex(Constraint):
193     def __call__(self, x):
194         x_sum = jnp.sum(x, axis=-1)
195         return jnp.all(x > 0, axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)
196 
197 
198 # TODO: Make types consistent
199 
200 boolean = _Boolean()
201 corr_cholesky = _CorrCholesky()
202 corr_matrix = _CorrMatrix()
203 dependent = _Dependent()
204 greater_than = _GreaterThan
205 less_than = _LessThan
206 integer_interval = _IntegerInterval
207 integer_greater_than = _IntegerGreaterThan
208 interval = _Interval
209 lower_cholesky = _LowerCholesky()
210 multinomial = _Multinomial
211 nonnegative_integer = _IntegerGreaterThan(0)
212 ordered_vector = _OrderedVector()
213 positive = _GreaterThan(0.)
214 positive_definite = _PositiveDefinite()
215 positive_integer = _IntegerGreaterThan(1)
216 real = _Real()
217 real_vector = _RealVector()
218 simplex = _Simplex()
219 unit_interval = _Interval(0., 1.)
220
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
 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

{"golden_diff": "diff --git a/numpyro/distributions/constraints.py b/numpyro/distributions/constraints.py\n--- a/numpyro/distributions/constraints.py\n+++ b/numpyro/distributions/constraints.py\n@@ -192,7 +192,7 @@\n class _Simplex(Constraint):\n     def __call__(self, x):\n         x_sum = jnp.sum(x, axis=-1)\n-        return jnp.all(x > 0, axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)\n+        return jnp.all(x >= 0, axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)\n \n \n # TODO: Make types consistent\n", "issue": "Possible error in the validation of a Categorical distribution\nI am getting  an error when I try to run the following code. The code just sample from a categorical distribution using the defined probabilities. \r\n\r\n```python\r\nimport numpyro\r\nimport numpyro.distributions as dist\r\nimport jax.numpy as jnp\r\nnumpyro.enable_validation(True)\r\n\r\n\r\ndef model():\r\n    probs = jnp.array([0.5, 0.5, 0.])\r\n    c = numpyro.sample('c', dist.Categorical(probs=probs))\r\n    return c\r\n\r\nwith numpyro.handlers.seed(rng_seed=54):\r\n    print(model())\r\n```\r\n\r\n```\r\nValueError                                Traceback (most recent call last)\r\n<ipython-input-1-fc7fe60e083b> in <module>\r\n     10 \r\n     11 with numpyro.handlers.seed(rng_seed=54):\r\n---> 12     print(model())\r\n\r\n<ipython-input-1-fc7fe60e083b> in model()\r\n      6 def model():\r\n      7     probs = jnp.array([0.5, 0.5, 0.])\r\n----> 8     c = numpyro.sample('c', dist.Categorical(probs=probs))\r\n      9     return c\r\n     10 \r\n\r\n~/miniconda3/envs/numpyro_test/lib/python3.8/site-packages/numpyro/distributions/discrete.py in Categorical(probs, logits, validate_args)\r\n    348 def Categorical(probs=None, logits=None, validate_args=None):\r\n    349     if probs is not None:\r\n--> 350         return CategoricalProbs(probs, validate_args=validate_args)\r\n    351     elif logits is not None:\r\n    352         return CategoricalLogits(logits, validate_args=validate_args)\r\n\r\n~/miniconda3/envs/numpyro_test/lib/python3.8/site-packages/numpyro/distributions/discrete.py in __init__(self, probs, validate_args)\r\n    265             raise ValueError(\"`probs` parameter must be at least one-dimensional.\")\r\n    266         self.probs = probs\r\n--> 267         super(CategoricalProbs, self).__init__(batch_shape=jnp.shape(self.probs)[:-1],\r\n    268                                                validate_args=validate_args)\r\n    269 \r\n\r\n~/miniconda3/envs/numpyro_test/lib/python3.8/site-packages/numpyro/distributions/distribution.py in __init__(self, batch_shape, event_shape, validate_args)\r\n    142                 if not_jax_tracer(is_valid):\r\n    143                     if not is_valid:\r\n--> 144                         raise ValueError(\"The parameter {} has invalid values\".format(param))\r\n    145         super(Distribution, self).__init__()\r\n    146 \r\n\r\nValueError: The parameter probs has invalid values\r\n```\r\n\r\nI 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.\r\n\r\n\r\nIf I write a similar code in Pyro with the validation enabled, I do not get an error.\r\n\r\n```python\r\nimport torch\r\nimport pyro\r\nimport pyro.distributions as dist\r\npyro.enable_validation(True)\r\npyro.set_rng_seed(54)\r\n\r\ndef model():\r\n    probs = torch.tensor([0.5, 0.5, 0.])\r\n    c = pyro.sample('c', dist.Categorical(probs=probs))\r\n    return c\r\n\r\nprint(model())\r\n```\r\n\r\nI am using Python 3.8.5, Pyro 1.4.0 and NumPyro 0.3.0 with Ubuntu. Happy to help with what I can.\n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\n# The implementation follows the design in PyTorch: torch.distributions.constraints.py\n#\n# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)\n# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)\n# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)\n# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n# Copyright (c) 2011-2013 NYU                      (Clement Farabet)\n# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)\n# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n#\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n\n\n__all__ = [\n    'boolean',\n    'corr_cholesky',\n    'corr_matrix',\n    'dependent',\n    'greater_than',\n    'integer_interval',\n    'integer_greater_than',\n    'interval',\n    'is_dependent',\n    'less_than',\n    'lower_cholesky',\n    'multinomial',\n    'nonnegative_integer',\n    'positive',\n    'positive_definite',\n    'positive_integer',\n    'real',\n    'real_vector',\n    'simplex',\n    'unit_interval',\n    'Constraint',\n]\n\nimport jax.numpy as jnp\n\n\nclass Constraint(object):\n    \"\"\"\n    Abstract base class for constraints.\n\n    A constraint object represents a region over which a variable is valid,\n    e.g. within which a variable can be optimized.\n    \"\"\"\n\n    def __call__(self, x):\n        raise NotImplementedError\n\n    def check(self, value):\n        \"\"\"\n        Returns a byte tensor of `sample_shape + batch_shape` indicating\n        whether each event in value satisfies this constraint.\n        \"\"\"\n        return self(value)\n\n\nclass _Boolean(Constraint):\n    def __call__(self, x):\n        return (x == 0) | (x == 1)\n\n\nclass _CorrCholesky(Constraint):\n    def __call__(self, x):\n        tril = jnp.tril(x)\n        lower_triangular = jnp.all(jnp.reshape(tril == x, x.shape[:-2] + (-1,)), axis=-1)\n        positive_diagonal = jnp.all(jnp.diagonal(x, axis1=-2, axis2=-1) > 0, axis=-1)\n        x_norm = jnp.linalg.norm(x, axis=-1)\n        unit_norm_row = jnp.all((x_norm <= 1) & (x_norm > 1 - 1e-6), axis=-1)\n        return lower_triangular & positive_diagonal & unit_norm_row\n\n\nclass _CorrMatrix(Constraint):\n    def __call__(self, x):\n        # check for symmetric\n        symmetric = jnp.all(jnp.all(x == jnp.swapaxes(x, -2, -1), axis=-1), axis=-1)\n        # check for the smallest eigenvalue is positive\n        positive = jnp.linalg.eigh(x)[0][..., 0] > 0\n        # check for diagonal equal to 1\n        unit_variance = jnp.all(jnp.abs(jnp.diagonal(x, axis1=-2, axis2=-1) - 1) < 1e-6, axis=-1)\n        return symmetric & positive & unit_variance\n\n\nclass _Dependent(Constraint):\n    def __call__(self, x):\n        raise ValueError('Cannot determine validity of dependent constraint')\n\n\ndef is_dependent(constraint):\n    return isinstance(constraint, _Dependent)\n\n\nclass _GreaterThan(Constraint):\n    def __init__(self, lower_bound):\n        self.lower_bound = lower_bound\n\n    def __call__(self, x):\n        return x > self.lower_bound\n\n\nclass _LessThan(Constraint):\n    def __init__(self, upper_bound):\n        self.upper_bound = upper_bound\n\n    def __call__(self, x):\n        return x < self.upper_bound\n\n\nclass _IntegerInterval(Constraint):\n    def __init__(self, lower_bound, upper_bound):\n        self.lower_bound = lower_bound\n        self.upper_bound = upper_bound\n\n    def __call__(self, x):\n        return (x >= self.lower_bound) & (x <= self.upper_bound) & (x == jnp.floor(x))\n\n\nclass _IntegerGreaterThan(Constraint):\n    def __init__(self, lower_bound):\n        self.lower_bound = lower_bound\n\n    def __call__(self, x):\n        return (x % 1 == 0) & (x >= self.lower_bound)\n\n\nclass _Interval(Constraint):\n    def __init__(self, lower_bound, upper_bound):\n        self.lower_bound = lower_bound\n        self.upper_bound = upper_bound\n\n    def __call__(self, x):\n        return (x > self.lower_bound) & (x < self.upper_bound)\n\n\nclass _LowerCholesky(Constraint):\n    def __call__(self, x):\n        tril = jnp.tril(x)\n        lower_triangular = jnp.all(jnp.reshape(tril == x, x.shape[:-2] + (-1,)), axis=-1)\n        positive_diagonal = jnp.all(jnp.diagonal(x, axis1=-2, axis2=-1) > 0, axis=-1)\n        return lower_triangular & positive_diagonal\n\n\nclass _Multinomial(Constraint):\n    def __init__(self, upper_bound):\n        self.upper_bound = upper_bound\n\n    def __call__(self, x):\n        return jnp.all(x >= 0, axis=-1) & (jnp.sum(x, -1) == self.upper_bound)\n\n\nclass _OrderedVector(Constraint):\n    def __call__(self, x):\n        return jnp.all(x[..., 1:] > x[..., :-1], axis=-1)\n\n\nclass _PositiveDefinite(Constraint):\n    def __call__(self, x):\n        # check for symmetric\n        symmetric = jnp.all(jnp.all(x == jnp.swapaxes(x, -2, -1), axis=-1), axis=-1)\n        # check for the smallest eigenvalue is positive\n        positive = jnp.linalg.eigh(x)[0][..., 0] > 0\n        return symmetric & positive\n\n\nclass _Real(Constraint):\n    def __call__(self, x):\n        return jnp.isfinite(x)\n\n\nclass _RealVector(Constraint):\n    def __call__(self, x):\n        return jnp.all(jnp.isfinite(x), axis=-1)\n\n\nclass _Simplex(Constraint):\n    def __call__(self, x):\n        x_sum = jnp.sum(x, axis=-1)\n        return jnp.all(x > 0, axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)\n\n\n# TODO: Make types consistent\n\nboolean = _Boolean()\ncorr_cholesky = _CorrCholesky()\ncorr_matrix = _CorrMatrix()\ndependent = _Dependent()\ngreater_than = _GreaterThan\nless_than = _LessThan\ninteger_interval = _IntegerInterval\ninteger_greater_than = _IntegerGreaterThan\ninterval = _Interval\nlower_cholesky = _LowerCholesky()\nmultinomial = _Multinomial\nnonnegative_integer = _IntegerGreaterThan(0)\nordered_vector = _OrderedVector()\npositive = _GreaterThan(0.)\npositive_definite = _PositiveDefinite()\npositive_integer = _IntegerGreaterThan(1)\nreal = _Real()\nreal_vector = _RealVector()\nsimplex = _Simplex()\nunit_interval = _Interval(0., 1.)\n", "path": "numpyro/distributions/constraints.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\n# The implementation follows the design in PyTorch: torch.distributions.constraints.py\n#\n# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)\n# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)\n# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)\n# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n# Copyright (c) 2011-2013 NYU                      (Clement Farabet)\n# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)\n# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n#\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n\n\n__all__ = [\n    'boolean',\n    'corr_cholesky',\n    'corr_matrix',\n    'dependent',\n    'greater_than',\n    'integer_interval',\n    'integer_greater_than',\n    'interval',\n    'is_dependent',\n    'less_than',\n    'lower_cholesky',\n    'multinomial',\n    'nonnegative_integer',\n    'positive',\n    'positive_definite',\n    'positive_integer',\n    'real',\n    'real_vector',\n    'simplex',\n    'unit_interval',\n    'Constraint',\n]\n\nimport jax.numpy as jnp\n\n\nclass Constraint(object):\n    \"\"\"\n    Abstract base class for constraints.\n\n    A constraint object represents a region over which a variable is valid,\n    e.g. within which a variable can be optimized.\n    \"\"\"\n\n    def __call__(self, x):\n        raise NotImplementedError\n\n    def check(self, value):\n        \"\"\"\n        Returns a byte tensor of `sample_shape + batch_shape` indicating\n        whether each event in value satisfies this constraint.\n        \"\"\"\n        return self(value)\n\n\nclass _Boolean(Constraint):\n    def __call__(self, x):\n        return (x == 0) | (x == 1)\n\n\nclass _CorrCholesky(Constraint):\n    def __call__(self, x):\n        tril = jnp.tril(x)\n        lower_triangular = jnp.all(jnp.reshape(tril == x, x.shape[:-2] + (-1,)), axis=-1)\n        positive_diagonal = jnp.all(jnp.diagonal(x, axis1=-2, axis2=-1) > 0, axis=-1)\n        x_norm = jnp.linalg.norm(x, axis=-1)\n        unit_norm_row = jnp.all((x_norm <= 1) & (x_norm > 1 - 1e-6), axis=-1)\n        return lower_triangular & positive_diagonal & unit_norm_row\n\n\nclass _CorrMatrix(Constraint):\n    def __call__(self, x):\n        # check for symmetric\n        symmetric = jnp.all(jnp.all(x == jnp.swapaxes(x, -2, -1), axis=-1), axis=-1)\n        # check for the smallest eigenvalue is positive\n        positive = jnp.linalg.eigh(x)[0][..., 0] > 0\n        # check for diagonal equal to 1\n        unit_variance = jnp.all(jnp.abs(jnp.diagonal(x, axis1=-2, axis2=-1) - 1) < 1e-6, axis=-1)\n        return symmetric & positive & unit_variance\n\n\nclass _Dependent(Constraint):\n    def __call__(self, x):\n        raise ValueError('Cannot determine validity of dependent constraint')\n\n\ndef is_dependent(constraint):\n    return isinstance(constraint, _Dependent)\n\n\nclass _GreaterThan(Constraint):\n    def __init__(self, lower_bound):\n        self.lower_bound = lower_bound\n\n    def __call__(self, x):\n        return x > self.lower_bound\n\n\nclass _LessThan(Constraint):\n    def __init__(self, upper_bound):\n        self.upper_bound = upper_bound\n\n    def __call__(self, x):\n        return x < self.upper_bound\n\n\nclass _IntegerInterval(Constraint):\n    def __init__(self, lower_bound, upper_bound):\n        self.lower_bound = lower_bound\n        self.upper_bound = upper_bound\n\n    def __call__(self, x):\n        return (x >= self.lower_bound) & (x <= self.upper_bound) & (x == jnp.floor(x))\n\n\nclass _IntegerGreaterThan(Constraint):\n    def __init__(self, lower_bound):\n        self.lower_bound = lower_bound\n\n    def __call__(self, x):\n        return (x % 1 == 0) & (x >= self.lower_bound)\n\n\nclass _Interval(Constraint):\n    def __init__(self, lower_bound, upper_bound):\n        self.lower_bound = lower_bound\n        self.upper_bound = upper_bound\n\n    def __call__(self, x):\n        return (x > self.lower_bound) & (x < self.upper_bound)\n\n\nclass _LowerCholesky(Constraint):\n    def __call__(self, x):\n        tril = jnp.tril(x)\n        lower_triangular = jnp.all(jnp.reshape(tril == x, x.shape[:-2] + (-1,)), axis=-1)\n        positive_diagonal = jnp.all(jnp.diagonal(x, axis1=-2, axis2=-1) > 0, axis=-1)\n        return lower_triangular & positive_diagonal\n\n\nclass _Multinomial(Constraint):\n    def __init__(self, upper_bound):\n        self.upper_bound = upper_bound\n\n    def __call__(self, x):\n        return jnp.all(x >= 0, axis=-1) & (jnp.sum(x, -1) == self.upper_bound)\n\n\nclass _OrderedVector(Constraint):\n    def __call__(self, x):\n        return jnp.all(x[..., 1:] > x[..., :-1], axis=-1)\n\n\nclass _PositiveDefinite(Constraint):\n    def __call__(self, x):\n        # check for symmetric\n        symmetric = jnp.all(jnp.all(x == jnp.swapaxes(x, -2, -1), axis=-1), axis=-1)\n        # check for the smallest eigenvalue is positive\n        positive = jnp.linalg.eigh(x)[0][..., 0] > 0\n        return symmetric & positive\n\n\nclass _Real(Constraint):\n    def __call__(self, x):\n        return jnp.isfinite(x)\n\n\nclass _RealVector(Constraint):\n    def __call__(self, x):\n        return jnp.all(jnp.isfinite(x), axis=-1)\n\n\nclass _Simplex(Constraint):\n    def __call__(self, x):\n        x_sum = jnp.sum(x, axis=-1)\n        return jnp.all(x >= 0, axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)\n\n\n# TODO: Make types consistent\n\nboolean = _Boolean()\ncorr_cholesky = _CorrCholesky()\ncorr_matrix = _CorrMatrix()\ndependent = _Dependent()\ngreater_than = _GreaterThan\nless_than = _LessThan\ninteger_interval = _IntegerInterval\ninteger_greater_than = _IntegerGreaterThan\ninterval = _Interval\nlower_cholesky = _LowerCholesky()\nmultinomial = _Multinomial\nnonnegative_integer = _IntegerGreaterThan(0)\nordered_vector = _OrderedVector()\npositive = _GreaterThan(0.)\npositive_definite = _PositiveDefinite()\npositive_integer = _IntegerGreaterThan(1)\nreal = _Real()\nreal_vector = _RealVector()\nsimplex = _Simplex()\nunit_interval = _Interval(0., 1.)\n", "path": "numpyro/distributions/constraints.py"}]}