We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Declare dependent dimensions
In Pyro, we can declare dependent dimensions through `.to_event()`. Currently, this is useful for correctly computing log_prob of a transformed distribution where `transform.event_dim>0` but `base_dist.event_shape=()`. In particular, we need this to declare `base_dist` in autoguide.

In Pyro, this is achieved by using Independent class. It seems unnecessary to me. Can we just manipulate `batch_shape`, `event_shape` and correctly compute the returned value for `log_prob`? I think we just need 1 line of code for each distribution
```
def log_prob(...):
    ...
    # event_ndim minus 1 for multivariate distributions
    return sum_rightmost(log_prob, len(self.event_shape))
```

If Independent class turns out to be useful, then we will implement it instead.

--- 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/distribution.py`
Content:
```
1 # The implementation follows the design in PyTorch: torch.distributions.distribution.py
2 #
3 # Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
4 # Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
5 # Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
6 # Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
7 # Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
8 # Copyright (c) 2011-2013 NYU                      (Clement Farabet)
9 # Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
10 # Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
11 # Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
12 #
13 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
14 # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
15 # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
16 # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
17 # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
18 # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
19 # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
20 # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
21 # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
22 # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
23 # POSSIBILITY OF SUCH DAMAGE.
24 
25 import jax.numpy as np
26 
27 from numpyro.distributions.constraints import Transform, is_dependent
28 from numpyro.distributions.util import lazy_property, sum_rightmost
29 
30 
31 class Distribution(object):
32     """
33     Base class for probability distributions in NumPyro. The design largely
34     follows from :mod:`torch.distributions`.
35 
36     :param batch_shape: The batch shape for the distribution. This designates
37         independent (possibly non-identical) dimensions of a sample from the
38         distribution. This is fixed for a distribution instance and is inferred
39         from the shape of the distribution parameters.
40     :param event_shape: The event shape for the distribution. This designates
41         the dependent dimensions of a sample from the distribution. These are
42         collapsed when we evaluate the log probability density of a batch of
43         samples using `.log_prob`.
44     :param validate_args: Whether to enable validation of distribution
45         parameters and arguments to `.log_prob` method.
46 
47     As an example:
48 
49     .. testsetup::
50 
51        import jax.numpy as np
52        import numpyro.distributions as dist
53 
54     .. doctest::
55 
56        >>> d = dist.Dirichlet(np.ones((2, 3, 4)))
57        >>> d.batch_shape
58        (2, 3)
59        >>> d.event_shape
60        (4,)
61     """
62     arg_constraints = {}
63     support = None
64     reparametrized_params = []
65     _validate_args = False
66 
67     def __init__(self, batch_shape=(), event_shape=(), validate_args=None):
68         self._batch_shape = batch_shape
69         self._event_shape = event_shape
70         if validate_args is not None:
71             self._validate_args = validate_args
72         if self._validate_args:
73             for param, constraint in self.arg_constraints.items():
74                 if param not in self.__dict__ and isinstance(getattr(type(self), param), lazy_property):
75                     continue
76                 if is_dependent(constraint):
77                     continue  # skip constraints that cannot be checked
78                 if not np.all(constraint(getattr(self, param))):
79                     raise ValueError("The parameter {} has invalid values".format(param))
80         super(Distribution, self).__init__()
81 
82     @property
83     def batch_shape(self):
84         """
85         Returns the shape over which the distribution parameters are batched.
86 
87         :return: batch shape of the distribution.
88         :rtype: tuple
89         """
90         return self._batch_shape
91 
92     @property
93     def event_shape(self):
94         """
95         Returns the shape of a single sample from the distribution without
96         batching.
97 
98         :return: event shape of the distribution.
99         :rtype: tuple
100         """
101         return self._event_shape
102 
103     def sample(self, key, sample_shape=()):
104         """
105         Returns a sample from the distribution having shape given by
106         `sample_shape + batch_shape + event_shape`. Note that when `sample_shape` is non-empty,
107         leading dimensions (of size `sample_shape`) of the returned sample will
108         be filled with iid draws from the distribution instance.
109 
110         :param jax.random.PRNGKey key: the rng key to be used for the distribution.
111         :param tuple sample_shape: the sample shape for the distribution.
112         :return: an array of shape `sample_shape + batch_shape + event_shape`
113         :rtype: numpy.ndarray
114         """
115         raise NotImplementedError
116 
117     def sample_with_intermediates(self, key, sample_shape=()):
118         """
119         Same as ``sample`` except that any intermediate computations are
120         returned (useful for `TransformedDistribution`).
121 
122         :param jax.random.PRNGKey key: the rng key to be used for the distribution.
123         :param tuple sample_shape: the sample shape for the distribution.
124         :return: an array of shape `sample_shape + batch_shape + event_shape`
125         :rtype: numpy.ndarray
126         """
127         return self.sample(key, sample_shape=sample_shape), []
128 
129     def transform_with_intermediates(self, base_value):
130         return base_value, []
131 
132     def log_prob(self, value):
133         """
134         Evaluates the log probability density for a batch of samples given by
135         `value`.
136 
137         :param value: A batch of samples from the distribution.
138         :return: an array with shape `value.shape[:-self.event_shape]`
139         :rtype: numpy.ndarray
140         """
141         raise NotImplementedError
142 
143     @property
144     def mean(self):
145         """
146         Mean of the distribution.
147         """
148         raise NotImplementedError
149 
150     @property
151     def variance(self):
152         """
153         Variance of the distribution.
154         """
155         raise NotImplementedError
156 
157     def _validate_sample(self, value):
158         if not np.all(self.support(value)):
159             raise ValueError('Invalid values provided to log prob method. '
160                              'The value argument must be within the support.')
161 
162     def __call__(self, *args, **kwargs):
163         key = kwargs.pop('random_state')
164         sample_intermediates = kwargs.pop('sample_intermediates', False)
165         if sample_intermediates:
166             return self.sample_with_intermediates(key, *args, **kwargs)
167         return self.sample(key, *args, **kwargs)
168 
169 
170 class TransformedDistribution(Distribution):
171     """
172     Returns a distribution instance obtained as a result of applying
173     a sequence of transforms to a base distribution. For an example,
174     see :class:`~numpyro.distributions.LogNormal` and
175     :class:`~numpyro.distributions.HalfNormal`.
176 
177     :param base_distribution: the base distribution over which to apply transforms.
178     :param transforms: a single transform or a list of transforms.
179     :param validate_args: Whether to enable validation of distribution
180         parameters and arguments to `.log_prob` method.
181     """
182     arg_constraints = {}
183 
184     def __init__(self, base_distribution, transforms, validate_args=None):
185         if isinstance(transforms, Transform):
186             transforms = [transforms, ]
187         elif isinstance(transforms, list):
188             if not all(isinstance(t, Transform) for t in transforms):
189                 raise ValueError("transforms must be a Transform or a list of Transforms")
190         else:
191             raise ValueError("transforms must be a Transform or list, but was {}".format(transforms))
192         # XXX: this logic will not be valid when IndependentDistribution is support;
193         # in that case, it is more involved to support Transform(Indep(Transform));
194         # however, we might not need to support such kind of distribution
195         # and should raise an error if base_distribution is an Indep one
196         if isinstance(base_distribution, TransformedDistribution):
197             self.base_dist = base_distribution.base_dist
198             self.transforms = base_distribution.transforms + transforms
199         else:
200             self.base_dist = base_distribution
201             self.transforms = transforms
202         shape = base_distribution.batch_shape + base_distribution.event_shape
203         event_dim = max([len(base_distribution.event_shape)] + [t.event_dim for t in transforms])
204         batch_shape = shape[:len(shape) - event_dim]
205         event_shape = shape[len(shape) - event_dim:]
206         super(TransformedDistribution, self).__init__(batch_shape, event_shape, validate_args=validate_args)
207 
208     @property
209     def support(self):
210         domain = self.base_dist.support
211         for t in self.transforms:
212             t.domain = domain
213             domain = t.codomain
214         return domain
215 
216     def sample(self, key, sample_shape=()):
217         x = self.base_dist.sample(key, sample_shape)
218         for transform in self.transforms:
219             x = transform(x)
220         return x
221 
222     def sample_with_intermediates(self, key, sample_shape=()):
223         base_value = self.base_dist.sample(key, sample_shape)
224         return self.transform_with_intermediates(base_value)
225 
226     def transform_with_intermediates(self, base_value):
227         x = base_value
228         intermediates = []
229         for transform in self.transforms:
230             x_tmp = x
231             x, t_inter = transform.call_with_intermediates(x)
232             intermediates.append([x_tmp, t_inter])
233         return x, intermediates
234 
235     def log_prob(self, value, intermediates=None):
236         if self._validate_args:
237             self._validate_sample(value)
238         if intermediates is not None:
239             if len(intermediates) != len(self.transforms):
240                 raise ValueError('Intermediates array has length = {}. Expected = {}.'
241                                  .format(len(intermediates), len(self.transforms)))
242         event_dim = len(self.event_shape)
243         log_prob = 0.0
244         y = value
245         for i, transform in enumerate(reversed(self.transforms)):
246             x = transform.inv(y) if intermediates is None else intermediates[-i - 1][0]
247             t_inter = None if intermediates is None else intermediates[-i - 1][1]
248             t_log_det = transform.log_abs_det_jacobian(x, y, t_inter)
249             log_prob = log_prob - sum_rightmost(t_log_det, event_dim - transform.event_dim)
250             y = x
251 
252         log_prob = log_prob + sum_rightmost(self.base_dist.log_prob(y),
253                                             event_dim - len(self.base_dist.event_shape))
254         return log_prob
255 
256     @property
257     def mean(self):
258         raise NotImplementedError
259 
260     @property
261     def variance(self):
262         raise NotImplementedError
263
```

Path: `numpyro/distributions/__init__.py`
Content:
```
1 from numpyro.distributions.continuous import (
2     Beta,
3     Cauchy,
4     Chi2,
5     Dirichlet,
6     Exponential,
7     Gamma,
8     GaussianRandomWalk,
9     HalfCauchy,
10     HalfNormal,
11     InverseGamma,
12     LKJCholesky,
13     LogNormal,
14     MultivariateNormal,
15     Normal,
16     Pareto,
17     StudentT,
18     TruncatedCauchy,
19     TruncatedNormal,
20     Uniform
21 )
22 from numpyro.distributions.discrete import (
23     Bernoulli,
24     BernoulliLogits,
25     BernoulliProbs,
26     Binomial,
27     BinomialLogits,
28     BinomialProbs,
29     Categorical,
30     CategoricalLogits,
31     CategoricalProbs,
32     Delta,
33     Multinomial,
34     MultinomialLogits,
35     MultinomialProbs,
36     Poisson,
37     PRNGIdentity
38 )
39 from numpyro.distributions.distribution import Distribution, TransformedDistribution
40 
41 __all__ = [
42     'Bernoulli',
43     'BernoulliLogits',
44     'BernoulliProbs',
45     'Beta',
46     'Binomial',
47     'BinomialLogits',
48     'BinomialProbs',
49     'Categorical',
50     'CategoricalLogits',
51     'CategoricalProbs',
52     'Cauchy',
53     'Chi2',
54     'Delta',
55     'Dirichlet',
56     'Distribution',
57     'Exponential',
58     'Gamma',
59     'GaussianRandomWalk',
60     'HalfCauchy',
61     'HalfNormal',
62     'InverseGamma',
63     'LKJCholesky',
64     'LogNormal',
65     'Multinomial',
66     'MultinomialLogits',
67     'MultinomialProbs',
68     'MultivariateNormal',
69     'Normal',
70     'Pareto',
71     'Poisson',
72     'PRNGIdentity',
73     'StudentT',
74     'TransformedDistribution',
75     'TruncatedCauchy',
76     'TruncatedNormal',
77     'Uniform',
78 ]
79
```
--- 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/__init__.py b/numpyro/distributions/__init__.py
--- a/numpyro/distributions/__init__.py
+++ b/numpyro/distributions/__init__.py
@@ -36,7 +36,7 @@
     Poisson,
     PRNGIdentity
 )
-from numpyro.distributions.distribution import Distribution, TransformedDistribution
+from numpyro.distributions.distribution import Distribution, Independent, TransformedDistribution
 
 __all__ = [
     'Bernoulli',
@@ -59,6 +59,7 @@
     'GaussianRandomWalk',
     'HalfCauchy',
     'HalfNormal',
+    'Independent',
     'InverseGamma',
     'LKJCholesky',
     'LogNormal',
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -166,6 +166,75 @@
             return self.sample_with_intermediates(key, *args, **kwargs)
         return self.sample(key, *args, **kwargs)
 
+    def to_event(self, reinterpreted_batch_ndims=None):
+        if reinterpreted_batch_ndims is None:
+            reinterpreted_batch_ndims = len(self.batch_shape)
+        return Independent(self, reinterpreted_batch_ndims)
+
+
+class Independent(Distribution):
+    """
+    Reinterprets batch dimensions of a distribution as event dims by shifting
+    the batch-event dim boundary further to the left.
+
+    From a practical standpoint, this is useful when changing the result of
+    :meth:`log_prob`. For example, a univariate Normal distribution can be
+    interpreted as a multivariate Normal with diagonal covariance::
+
+    .. testsetup::
+
+       import numpyro.distributions as dist
+
+    .. doctest::
+
+        >>> normal = dist.Normal(np.zeros(3), np.ones(3))
+        >>> [normal.batch_shape, normal.event_shape]
+        [torch.Size((3,)), torch.Size(())]
+        >>> diag_normal = Independent(normal, 1)
+        >>> [diag_normal.batch_shape, diag_normal.event_shape]
+        [torch.Size(()), torch.Size((3,))]
+
+    :param numpyro.distribution.Distribution base_distribution: a distribution instance.
+    :param int reinterpreted_batch_ndims: the number of batch dims to reinterpret as event dims.
+    """
+    arg_constraints = {}
+
+    def __init__(self, base_dist, reinterpreted_batch_ndims, validate_args=None):
+        if reinterpreted_batch_ndims > len(base_dist.batch_shape):
+            raise ValueError("Expected reinterpreted_batch_ndims <= len(base_distribution.batch_shape), "
+                             "actual {} vs {}".format(reinterpreted_batch_ndims,
+                                                      len(base_dist.batch_shape)))
+        shape = base_dist.batch_shape + base_dist.event_shape
+        event_dim = reinterpreted_batch_ndims + len(base_dist.event_shape)
+        batch_shape = shape[:len(shape) - event_dim]
+        event_shape = shape[len(shape) - event_dim:]
+        self.base_dist = base_dist
+        self.reinterpreted_batch_ndims = reinterpreted_batch_ndims
+        super(Independent, self).__init__(batch_shape, event_shape, validate_args=validate_args)
+
+    @property
+    def support(self):
+        return self.base_dist.support
+
+    @property
+    def reparameterized_params(self):
+        return self.base_dist.reparameterized_params
+
+    @property
+    def mean(self):
+        return self.base_dist.mean
+
+    @property
+    def variance(self):
+        return self.base_dist.variance
+
+    def sample(self, key, sample_shape=()):
+        return self.base_dist.sample(key, sample_shape=sample_shape)
+
+    def log_prob(self, value):
+        log_prob = self.base_dist.log_prob(value)
+        return sum_rightmost(log_prob, self.reinterpreted_batch_ndims)
+
 
 class TransformedDistribution(Distribution):
     """

{"golden_diff": "diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py\n--- a/numpyro/distributions/__init__.py\n+++ b/numpyro/distributions/__init__.py\n@@ -36,7 +36,7 @@\n     Poisson,\n     PRNGIdentity\n )\n-from numpyro.distributions.distribution import Distribution, TransformedDistribution\n+from numpyro.distributions.distribution import Distribution, Independent, TransformedDistribution\n \n __all__ = [\n     'Bernoulli',\n@@ -59,6 +59,7 @@\n     'GaussianRandomWalk',\n     'HalfCauchy',\n     'HalfNormal',\n+    'Independent',\n     'InverseGamma',\n     'LKJCholesky',\n     'LogNormal',\ndiff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py\n--- a/numpyro/distributions/distribution.py\n+++ b/numpyro/distributions/distribution.py\n@@ -166,6 +166,75 @@\n             return self.sample_with_intermediates(key, *args, **kwargs)\n         return self.sample(key, *args, **kwargs)\n \n+    def to_event(self, reinterpreted_batch_ndims=None):\n+        if reinterpreted_batch_ndims is None:\n+            reinterpreted_batch_ndims = len(self.batch_shape)\n+        return Independent(self, reinterpreted_batch_ndims)\n+\n+\n+class Independent(Distribution):\n+    \"\"\"\n+    Reinterprets batch dimensions of a distribution as event dims by shifting\n+    the batch-event dim boundary further to the left.\n+\n+    From a practical standpoint, this is useful when changing the result of\n+    :meth:`log_prob`. For example, a univariate Normal distribution can be\n+    interpreted as a multivariate Normal with diagonal covariance::\n+\n+    .. testsetup::\n+\n+       import numpyro.distributions as dist\n+\n+    .. doctest::\n+\n+        >>> normal = dist.Normal(np.zeros(3), np.ones(3))\n+        >>> [normal.batch_shape, normal.event_shape]\n+        [torch.Size((3,)), torch.Size(())]\n+        >>> diag_normal = Independent(normal, 1)\n+        >>> [diag_normal.batch_shape, diag_normal.event_shape]\n+        [torch.Size(()), torch.Size((3,))]\n+\n+    :param numpyro.distribution.Distribution base_distribution: a distribution instance.\n+    :param int reinterpreted_batch_ndims: the number of batch dims to reinterpret as event dims.\n+    \"\"\"\n+    arg_constraints = {}\n+\n+    def __init__(self, base_dist, reinterpreted_batch_ndims, validate_args=None):\n+        if reinterpreted_batch_ndims > len(base_dist.batch_shape):\n+            raise ValueError(\"Expected reinterpreted_batch_ndims <= len(base_distribution.batch_shape), \"\n+                             \"actual {} vs {}\".format(reinterpreted_batch_ndims,\n+                                                      len(base_dist.batch_shape)))\n+        shape = base_dist.batch_shape + base_dist.event_shape\n+        event_dim = reinterpreted_batch_ndims + len(base_dist.event_shape)\n+        batch_shape = shape[:len(shape) - event_dim]\n+        event_shape = shape[len(shape) - event_dim:]\n+        self.base_dist = base_dist\n+        self.reinterpreted_batch_ndims = reinterpreted_batch_ndims\n+        super(Independent, self).__init__(batch_shape, event_shape, validate_args=validate_args)\n+\n+    @property\n+    def support(self):\n+        return self.base_dist.support\n+\n+    @property\n+    def reparameterized_params(self):\n+        return self.base_dist.reparameterized_params\n+\n+    @property\n+    def mean(self):\n+        return self.base_dist.mean\n+\n+    @property\n+    def variance(self):\n+        return self.base_dist.variance\n+\n+    def sample(self, key, sample_shape=()):\n+        return self.base_dist.sample(key, sample_shape=sample_shape)\n+\n+    def log_prob(self, value):\n+        log_prob = self.base_dist.log_prob(value)\n+        return sum_rightmost(log_prob, self.reinterpreted_batch_ndims)\n+\n \n class TransformedDistribution(Distribution):\n     \"\"\"\n", "issue": "Declare dependent dimensions\nIn Pyro, we can declare dependent dimensions through `.to_event()`. Currently, this is useful for correctly computing log_prob of a transformed distribution where `transform.event_dim>0` but `base_dist.event_shape=()`. In particular, we need this to declare `base_dist` in autoguide.\r\n\r\nIn Pyro, this is achieved by using Independent class. It seems unnecessary to me. Can we just manipulate `batch_shape`, `event_shape` and correctly compute the returned value for `log_prob`? I think we just need 1 line of code for each distribution\r\n```\r\ndef log_prob(...):\r\n    ...\r\n    # event_ndim minus 1 for multivariate distributions\r\n    return sum_rightmost(log_prob, len(self.event_shape))\r\n```\r\n\r\nIf Independent class turns out to be useful, then we will implement it instead.\n", "before_files": [{"content": "# The implementation follows the design in PyTorch: torch.distributions.distribution.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\nimport jax.numpy as np\n\nfrom numpyro.distributions.constraints import Transform, is_dependent\nfrom numpyro.distributions.util import lazy_property, sum_rightmost\n\n\nclass Distribution(object):\n    \"\"\"\n    Base class for probability distributions in NumPyro. The design largely\n    follows from :mod:`torch.distributions`.\n\n    :param batch_shape: The batch shape for the distribution. This designates\n        independent (possibly non-identical) dimensions of a sample from the\n        distribution. This is fixed for a distribution instance and is inferred\n        from the shape of the distribution parameters.\n    :param event_shape: The event shape for the distribution. This designates\n        the dependent dimensions of a sample from the distribution. These are\n        collapsed when we evaluate the log probability density of a batch of\n        samples using `.log_prob`.\n    :param validate_args: Whether to enable validation of distribution\n        parameters and arguments to `.log_prob` method.\n\n    As an example:\n\n    .. testsetup::\n\n       import jax.numpy as np\n       import numpyro.distributions as dist\n\n    .. doctest::\n\n       >>> d = dist.Dirichlet(np.ones((2, 3, 4)))\n       >>> d.batch_shape\n       (2, 3)\n       >>> d.event_shape\n       (4,)\n    \"\"\"\n    arg_constraints = {}\n    support = None\n    reparametrized_params = []\n    _validate_args = False\n\n    def __init__(self, batch_shape=(), event_shape=(), validate_args=None):\n        self._batch_shape = batch_shape\n        self._event_shape = event_shape\n        if validate_args is not None:\n            self._validate_args = validate_args\n        if self._validate_args:\n            for param, constraint in self.arg_constraints.items():\n                if param not in self.__dict__ and isinstance(getattr(type(self), param), lazy_property):\n                    continue\n                if is_dependent(constraint):\n                    continue  # skip constraints that cannot be checked\n                if not np.all(constraint(getattr(self, param))):\n                    raise ValueError(\"The parameter {} has invalid values\".format(param))\n        super(Distribution, self).__init__()\n\n    @property\n    def batch_shape(self):\n        \"\"\"\n        Returns the shape over which the distribution parameters are batched.\n\n        :return: batch shape of the distribution.\n        :rtype: tuple\n        \"\"\"\n        return self._batch_shape\n\n    @property\n    def event_shape(self):\n        \"\"\"\n        Returns the shape of a single sample from the distribution without\n        batching.\n\n        :return: event shape of the distribution.\n        :rtype: tuple\n        \"\"\"\n        return self._event_shape\n\n    def sample(self, key, sample_shape=()):\n        \"\"\"\n        Returns a sample from the distribution having shape given by\n        `sample_shape + batch_shape + event_shape`. Note that when `sample_shape` is non-empty,\n        leading dimensions (of size `sample_shape`) of the returned sample will\n        be filled with iid draws from the distribution instance.\n\n        :param jax.random.PRNGKey key: the rng key to be used for the distribution.\n        :param tuple sample_shape: the sample shape for the distribution.\n        :return: an array of shape `sample_shape + batch_shape + event_shape`\n        :rtype: numpy.ndarray\n        \"\"\"\n        raise NotImplementedError\n\n    def sample_with_intermediates(self, key, sample_shape=()):\n        \"\"\"\n        Same as ``sample`` except that any intermediate computations are\n        returned (useful for `TransformedDistribution`).\n\n        :param jax.random.PRNGKey key: the rng key to be used for the distribution.\n        :param tuple sample_shape: the sample shape for the distribution.\n        :return: an array of shape `sample_shape + batch_shape + event_shape`\n        :rtype: numpy.ndarray\n        \"\"\"\n        return self.sample(key, sample_shape=sample_shape), []\n\n    def transform_with_intermediates(self, base_value):\n        return base_value, []\n\n    def log_prob(self, value):\n        \"\"\"\n        Evaluates the log probability density for a batch of samples given by\n        `value`.\n\n        :param value: A batch of samples from the distribution.\n        :return: an array with shape `value.shape[:-self.event_shape]`\n        :rtype: numpy.ndarray\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    def mean(self):\n        \"\"\"\n        Mean of the distribution.\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    def variance(self):\n        \"\"\"\n        Variance of the distribution.\n        \"\"\"\n        raise NotImplementedError\n\n    def _validate_sample(self, value):\n        if not np.all(self.support(value)):\n            raise ValueError('Invalid values provided to log prob method. '\n                             'The value argument must be within the support.')\n\n    def __call__(self, *args, **kwargs):\n        key = kwargs.pop('random_state')\n        sample_intermediates = kwargs.pop('sample_intermediates', False)\n        if sample_intermediates:\n            return self.sample_with_intermediates(key, *args, **kwargs)\n        return self.sample(key, *args, **kwargs)\n\n\nclass TransformedDistribution(Distribution):\n    \"\"\"\n    Returns a distribution instance obtained as a result of applying\n    a sequence of transforms to a base distribution. For an example,\n    see :class:`~numpyro.distributions.LogNormal` and\n    :class:`~numpyro.distributions.HalfNormal`.\n\n    :param base_distribution: the base distribution over which to apply transforms.\n    :param transforms: a single transform or a list of transforms.\n    :param validate_args: Whether to enable validation of distribution\n        parameters and arguments to `.log_prob` method.\n    \"\"\"\n    arg_constraints = {}\n\n    def __init__(self, base_distribution, transforms, validate_args=None):\n        if isinstance(transforms, Transform):\n            transforms = [transforms, ]\n        elif isinstance(transforms, list):\n            if not all(isinstance(t, Transform) for t in transforms):\n                raise ValueError(\"transforms must be a Transform or a list of Transforms\")\n        else:\n            raise ValueError(\"transforms must be a Transform or list, but was {}\".format(transforms))\n        # XXX: this logic will not be valid when IndependentDistribution is support;\n        # in that case, it is more involved to support Transform(Indep(Transform));\n        # however, we might not need to support such kind of distribution\n        # and should raise an error if base_distribution is an Indep one\n        if isinstance(base_distribution, TransformedDistribution):\n            self.base_dist = base_distribution.base_dist\n            self.transforms = base_distribution.transforms + transforms\n        else:\n            self.base_dist = base_distribution\n            self.transforms = transforms\n        shape = base_distribution.batch_shape + base_distribution.event_shape\n        event_dim = max([len(base_distribution.event_shape)] + [t.event_dim for t in transforms])\n        batch_shape = shape[:len(shape) - event_dim]\n        event_shape = shape[len(shape) - event_dim:]\n        super(TransformedDistribution, self).__init__(batch_shape, event_shape, validate_args=validate_args)\n\n    @property\n    def support(self):\n        domain = self.base_dist.support\n        for t in self.transforms:\n            t.domain = domain\n            domain = t.codomain\n        return domain\n\n    def sample(self, key, sample_shape=()):\n        x = self.base_dist.sample(key, sample_shape)\n        for transform in self.transforms:\n            x = transform(x)\n        return x\n\n    def sample_with_intermediates(self, key, sample_shape=()):\n        base_value = self.base_dist.sample(key, sample_shape)\n        return self.transform_with_intermediates(base_value)\n\n    def transform_with_intermediates(self, base_value):\n        x = base_value\n        intermediates = []\n        for transform in self.transforms:\n            x_tmp = x\n            x, t_inter = transform.call_with_intermediates(x)\n            intermediates.append([x_tmp, t_inter])\n        return x, intermediates\n\n    def log_prob(self, value, intermediates=None):\n        if self._validate_args:\n            self._validate_sample(value)\n        if intermediates is not None:\n            if len(intermediates) != len(self.transforms):\n                raise ValueError('Intermediates array has length = {}. Expected = {}.'\n                                 .format(len(intermediates), len(self.transforms)))\n        event_dim = len(self.event_shape)\n        log_prob = 0.0\n        y = value\n        for i, transform in enumerate(reversed(self.transforms)):\n            x = transform.inv(y) if intermediates is None else intermediates[-i - 1][0]\n            t_inter = None if intermediates is None else intermediates[-i - 1][1]\n            t_log_det = transform.log_abs_det_jacobian(x, y, t_inter)\n            log_prob = log_prob - sum_rightmost(t_log_det, event_dim - transform.event_dim)\n            y = x\n\n        log_prob = log_prob + sum_rightmost(self.base_dist.log_prob(y),\n                                            event_dim - len(self.base_dist.event_shape))\n        return log_prob\n\n    @property\n    def mean(self):\n        raise NotImplementedError\n\n    @property\n    def variance(self):\n        raise NotImplementedError\n", "path": "numpyro/distributions/distribution.py"}, {"content": "from numpyro.distributions.continuous import (\n    Beta,\n    Cauchy,\n    Chi2,\n    Dirichlet,\n    Exponential,\n    Gamma,\n    GaussianRandomWalk,\n    HalfCauchy,\n    HalfNormal,\n    InverseGamma,\n    LKJCholesky,\n    LogNormal,\n    MultivariateNormal,\n    Normal,\n    Pareto,\n    StudentT,\n    TruncatedCauchy,\n    TruncatedNormal,\n    Uniform\n)\nfrom numpyro.distributions.discrete import (\n    Bernoulli,\n    BernoulliLogits,\n    BernoulliProbs,\n    Binomial,\n    BinomialLogits,\n    BinomialProbs,\n    Categorical,\n    CategoricalLogits,\n    CategoricalProbs,\n    Delta,\n    Multinomial,\n    MultinomialLogits,\n    MultinomialProbs,\n    Poisson,\n    PRNGIdentity\n)\nfrom numpyro.distributions.distribution import Distribution, TransformedDistribution\n\n__all__ = [\n    'Bernoulli',\n    'BernoulliLogits',\n    'BernoulliProbs',\n    'Beta',\n    'Binomial',\n    'BinomialLogits',\n    'BinomialProbs',\n    'Categorical',\n    'CategoricalLogits',\n    'CategoricalProbs',\n    'Cauchy',\n    'Chi2',\n    'Delta',\n    'Dirichlet',\n    'Distribution',\n    'Exponential',\n    'Gamma',\n    'GaussianRandomWalk',\n    'HalfCauchy',\n    'HalfNormal',\n    'InverseGamma',\n    'LKJCholesky',\n    'LogNormal',\n    'Multinomial',\n    'MultinomialLogits',\n    'MultinomialProbs',\n    'MultivariateNormal',\n    'Normal',\n    'Pareto',\n    'Poisson',\n    'PRNGIdentity',\n    'StudentT',\n    'TransformedDistribution',\n    'TruncatedCauchy',\n    'TruncatedNormal',\n    'Uniform',\n]\n", "path": "numpyro/distributions/__init__.py"}], "after_files": [{"content": "# The implementation follows the design in PyTorch: torch.distributions.distribution.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\nimport jax.numpy as np\n\nfrom numpyro.distributions.constraints import Transform, is_dependent\nfrom numpyro.distributions.util import lazy_property, sum_rightmost\n\n\nclass Distribution(object):\n    \"\"\"\n    Base class for probability distributions in NumPyro. The design largely\n    follows from :mod:`torch.distributions`.\n\n    :param batch_shape: The batch shape for the distribution. This designates\n        independent (possibly non-identical) dimensions of a sample from the\n        distribution. This is fixed for a distribution instance and is inferred\n        from the shape of the distribution parameters.\n    :param event_shape: The event shape for the distribution. This designates\n        the dependent dimensions of a sample from the distribution. These are\n        collapsed when we evaluate the log probability density of a batch of\n        samples using `.log_prob`.\n    :param validate_args: Whether to enable validation of distribution\n        parameters and arguments to `.log_prob` method.\n\n    As an example:\n\n    .. testsetup::\n\n       import jax.numpy as np\n       import numpyro.distributions as dist\n\n    .. doctest::\n\n       >>> d = dist.Dirichlet(np.ones((2, 3, 4)))\n       >>> d.batch_shape\n       (2, 3)\n       >>> d.event_shape\n       (4,)\n    \"\"\"\n    arg_constraints = {}\n    support = None\n    reparametrized_params = []\n    _validate_args = False\n\n    def __init__(self, batch_shape=(), event_shape=(), validate_args=None):\n        self._batch_shape = batch_shape\n        self._event_shape = event_shape\n        if validate_args is not None:\n            self._validate_args = validate_args\n        if self._validate_args:\n            for param, constraint in self.arg_constraints.items():\n                if param not in self.__dict__ and isinstance(getattr(type(self), param), lazy_property):\n                    continue\n                if is_dependent(constraint):\n                    continue  # skip constraints that cannot be checked\n                if not np.all(constraint(getattr(self, param))):\n                    raise ValueError(\"The parameter {} has invalid values\".format(param))\n        super(Distribution, self).__init__()\n\n    @property\n    def batch_shape(self):\n        \"\"\"\n        Returns the shape over which the distribution parameters are batched.\n\n        :return: batch shape of the distribution.\n        :rtype: tuple\n        \"\"\"\n        return self._batch_shape\n\n    @property\n    def event_shape(self):\n        \"\"\"\n        Returns the shape of a single sample from the distribution without\n        batching.\n\n        :return: event shape of the distribution.\n        :rtype: tuple\n        \"\"\"\n        return self._event_shape\n\n    def sample(self, key, sample_shape=()):\n        \"\"\"\n        Returns a sample from the distribution having shape given by\n        `sample_shape + batch_shape + event_shape`. Note that when `sample_shape` is non-empty,\n        leading dimensions (of size `sample_shape`) of the returned sample will\n        be filled with iid draws from the distribution instance.\n\n        :param jax.random.PRNGKey key: the rng key to be used for the distribution.\n        :param tuple sample_shape: the sample shape for the distribution.\n        :return: an array of shape `sample_shape + batch_shape + event_shape`\n        :rtype: numpy.ndarray\n        \"\"\"\n        raise NotImplementedError\n\n    def sample_with_intermediates(self, key, sample_shape=()):\n        \"\"\"\n        Same as ``sample`` except that any intermediate computations are\n        returned (useful for `TransformedDistribution`).\n\n        :param jax.random.PRNGKey key: the rng key to be used for the distribution.\n        :param tuple sample_shape: the sample shape for the distribution.\n        :return: an array of shape `sample_shape + batch_shape + event_shape`\n        :rtype: numpy.ndarray\n        \"\"\"\n        return self.sample(key, sample_shape=sample_shape), []\n\n    def transform_with_intermediates(self, base_value):\n        return base_value, []\n\n    def log_prob(self, value):\n        \"\"\"\n        Evaluates the log probability density for a batch of samples given by\n        `value`.\n\n        :param value: A batch of samples from the distribution.\n        :return: an array with shape `value.shape[:-self.event_shape]`\n        :rtype: numpy.ndarray\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    def mean(self):\n        \"\"\"\n        Mean of the distribution.\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    def variance(self):\n        \"\"\"\n        Variance of the distribution.\n        \"\"\"\n        raise NotImplementedError\n\n    def _validate_sample(self, value):\n        if not np.all(self.support(value)):\n            raise ValueError('Invalid values provided to log prob method. '\n                             'The value argument must be within the support.')\n\n    def __call__(self, *args, **kwargs):\n        key = kwargs.pop('random_state')\n        sample_intermediates = kwargs.pop('sample_intermediates', False)\n        if sample_intermediates:\n            return self.sample_with_intermediates(key, *args, **kwargs)\n        return self.sample(key, *args, **kwargs)\n\n    def to_event(self, reinterpreted_batch_ndims=None):\n        if reinterpreted_batch_ndims is None:\n            reinterpreted_batch_ndims = len(self.batch_shape)\n        return Independent(self, reinterpreted_batch_ndims)\n\n\nclass Independent(Distribution):\n    \"\"\"\n    Reinterprets batch dimensions of a distribution as event dims by shifting\n    the batch-event dim boundary further to the left.\n\n    From a practical standpoint, this is useful when changing the result of\n    :meth:`log_prob`. For example, a univariate Normal distribution can be\n    interpreted as a multivariate Normal with diagonal covariance::\n\n    .. testsetup::\n\n       import numpyro.distributions as dist\n\n    .. doctest::\n\n        >>> normal = dist.Normal(np.zeros(3), np.ones(3))\n        >>> [normal.batch_shape, normal.event_shape]\n        [torch.Size((3,)), torch.Size(())]\n        >>> diag_normal = Independent(normal, 1)\n        >>> [diag_normal.batch_shape, diag_normal.event_shape]\n        [torch.Size(()), torch.Size((3,))]\n\n    :param numpyro.distribution.Distribution base_distribution: a distribution instance.\n    :param int reinterpreted_batch_ndims: the number of batch dims to reinterpret as event dims.\n    \"\"\"\n    arg_constraints = {}\n\n    def __init__(self, base_dist, reinterpreted_batch_ndims, validate_args=None):\n        if reinterpreted_batch_ndims > len(base_dist.batch_shape):\n            raise ValueError(\"Expected reinterpreted_batch_ndims <= len(base_distribution.batch_shape), \"\n                             \"actual {} vs {}\".format(reinterpreted_batch_ndims,\n                                                      len(base_dist.batch_shape)))\n        shape = base_dist.batch_shape + base_dist.event_shape\n        event_dim = reinterpreted_batch_ndims + len(base_dist.event_shape)\n        batch_shape = shape[:len(shape) - event_dim]\n        event_shape = shape[len(shape) - event_dim:]\n        self.base_dist = base_dist\n        self.reinterpreted_batch_ndims = reinterpreted_batch_ndims\n        super(Independent, self).__init__(batch_shape, event_shape, validate_args=validate_args)\n\n    @property\n    def support(self):\n        return self.base_dist.support\n\n    @property\n    def reparameterized_params(self):\n        return self.base_dist.reparameterized_params\n\n    @property\n    def mean(self):\n        return self.base_dist.mean\n\n    @property\n    def variance(self):\n        return self.base_dist.variance\n\n    def sample(self, key, sample_shape=()):\n        return self.base_dist.sample(key, sample_shape=sample_shape)\n\n    def log_prob(self, value):\n        log_prob = self.base_dist.log_prob(value)\n        return sum_rightmost(log_prob, self.reinterpreted_batch_ndims)\n\n\nclass TransformedDistribution(Distribution):\n    \"\"\"\n    Returns a distribution instance obtained as a result of applying\n    a sequence of transforms to a base distribution. For an example,\n    see :class:`~numpyro.distributions.LogNormal` and\n    :class:`~numpyro.distributions.HalfNormal`.\n\n    :param base_distribution: the base distribution over which to apply transforms.\n    :param transforms: a single transform or a list of transforms.\n    :param validate_args: Whether to enable validation of distribution\n        parameters and arguments to `.log_prob` method.\n    \"\"\"\n    arg_constraints = {}\n\n    def __init__(self, base_distribution, transforms, validate_args=None):\n        if isinstance(transforms, Transform):\n            transforms = [transforms, ]\n        elif isinstance(transforms, list):\n            if not all(isinstance(t, Transform) for t in transforms):\n                raise ValueError(\"transforms must be a Transform or a list of Transforms\")\n        else:\n            raise ValueError(\"transforms must be a Transform or list, but was {}\".format(transforms))\n        # XXX: this logic will not be valid when IndependentDistribution is support;\n        # in that case, it is more involved to support Transform(Indep(Transform));\n        # however, we might not need to support such kind of distribution\n        # and should raise an error if base_distribution is an Indep one\n        if isinstance(base_distribution, TransformedDistribution):\n            self.base_dist = base_distribution.base_dist\n            self.transforms = base_distribution.transforms + transforms\n        else:\n            self.base_dist = base_distribution\n            self.transforms = transforms\n        shape = base_distribution.batch_shape + base_distribution.event_shape\n        event_dim = max([len(base_distribution.event_shape)] + [t.event_dim for t in transforms])\n        batch_shape = shape[:len(shape) - event_dim]\n        event_shape = shape[len(shape) - event_dim:]\n        super(TransformedDistribution, self).__init__(batch_shape, event_shape, validate_args=validate_args)\n\n    @property\n    def support(self):\n        domain = self.base_dist.support\n        for t in self.transforms:\n            t.domain = domain\n            domain = t.codomain\n        return domain\n\n    def sample(self, key, sample_shape=()):\n        x = self.base_dist.sample(key, sample_shape)\n        for transform in self.transforms:\n            x = transform(x)\n        return x\n\n    def sample_with_intermediates(self, key, sample_shape=()):\n        base_value = self.base_dist.sample(key, sample_shape)\n        return self.transform_with_intermediates(base_value)\n\n    def transform_with_intermediates(self, base_value):\n        x = base_value\n        intermediates = []\n        for transform in self.transforms:\n            x_tmp = x\n            x, t_inter = transform.call_with_intermediates(x)\n            intermediates.append([x_tmp, t_inter])\n        return x, intermediates\n\n    def log_prob(self, value, intermediates=None):\n        if self._validate_args:\n            self._validate_sample(value)\n        if intermediates is not None:\n            if len(intermediates) != len(self.transforms):\n                raise ValueError('Intermediates array has length = {}. Expected = {}.'\n                                 .format(len(intermediates), len(self.transforms)))\n        event_dim = len(self.event_shape)\n        log_prob = 0.0\n        y = value\n        for i, transform in enumerate(reversed(self.transforms)):\n            x = transform.inv(y) if intermediates is None else intermediates[-i - 1][0]\n            t_inter = None if intermediates is None else intermediates[-i - 1][1]\n            t_log_det = transform.log_abs_det_jacobian(x, y, t_inter)\n            log_prob = log_prob - sum_rightmost(t_log_det, event_dim - transform.event_dim)\n            y = x\n\n        log_prob = log_prob + sum_rightmost(self.base_dist.log_prob(y),\n                                            event_dim - len(self.base_dist.event_shape))\n        return log_prob\n\n    @property\n    def mean(self):\n        raise NotImplementedError\n\n    @property\n    def variance(self):\n        raise NotImplementedError\n", "path": "numpyro/distributions/distribution.py"}, {"content": "from numpyro.distributions.continuous import (\n    Beta,\n    Cauchy,\n    Chi2,\n    Dirichlet,\n    Exponential,\n    Gamma,\n    GaussianRandomWalk,\n    HalfCauchy,\n    HalfNormal,\n    InverseGamma,\n    LKJCholesky,\n    LogNormal,\n    MultivariateNormal,\n    Normal,\n    Pareto,\n    StudentT,\n    TruncatedCauchy,\n    TruncatedNormal,\n    Uniform\n)\nfrom numpyro.distributions.discrete import (\n    Bernoulli,\n    BernoulliLogits,\n    BernoulliProbs,\n    Binomial,\n    BinomialLogits,\n    BinomialProbs,\n    Categorical,\n    CategoricalLogits,\n    CategoricalProbs,\n    Delta,\n    Multinomial,\n    MultinomialLogits,\n    MultinomialProbs,\n    Poisson,\n    PRNGIdentity\n)\nfrom numpyro.distributions.distribution import Distribution, Independent, TransformedDistribution\n\n__all__ = [\n    'Bernoulli',\n    'BernoulliLogits',\n    'BernoulliProbs',\n    'Beta',\n    'Binomial',\n    'BinomialLogits',\n    'BinomialProbs',\n    'Categorical',\n    'CategoricalLogits',\n    'CategoricalProbs',\n    'Cauchy',\n    'Chi2',\n    'Delta',\n    'Dirichlet',\n    'Distribution',\n    'Exponential',\n    'Gamma',\n    'GaussianRandomWalk',\n    'HalfCauchy',\n    'HalfNormal',\n    'Independent',\n    'InverseGamma',\n    'LKJCholesky',\n    'LogNormal',\n    'Multinomial',\n    'MultinomialLogits',\n    'MultinomialProbs',\n    'MultivariateNormal',\n    'Normal',\n    'Pareto',\n    'Poisson',\n    'PRNGIdentity',\n    'StudentT',\n    'TransformedDistribution',\n    'TruncatedCauchy',\n    'TruncatedNormal',\n    'Uniform',\n]\n", "path": "numpyro/distributions/__init__.py"}]}