We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[bug] Samples from LKJCorrCholesky have the wrong shape
### Issue Description
The batch and sample shapes for `LKJCorrCholesky` are swapped when sampling.

### Environment
OS: Ubuntu 16.04
Python: 3.7.2
Torch version: 1.6.0
Pyro version: 1.4.0

### Code Snippet
Minimal example:

```python
>>> dist = pyro.distributions.LKJCorrCholesky(d=3, eta=torch.ones(())).expand(12)
>>> # batch shape and event shape are as you'd expect
>>> dist.batch_shape
torch.Size([12])
>>> dist.event_shape
torch.Size([3, 3])
>>> # samples have correct shape when sample_shape=()
>>> dist.shape(())
torch.Size([12, 3, 3])
>>> dist.sample().shape
torch.Size([12, 3, 3])
>>> # samples have the wrong shape when sample_shape is non-unit
>>> dist.shape((4,))  # as expected
torch.Size([4, 12, 3, 3])
>>> dist.sample((4,)).shape
torch.Size([12, 4, 3, 3])
```

I think this line is the culprit: https://github.com/pyro-ppl/pyro/blob/7c2c22c10dffda8a33ffbd593cc8d58819959e40/pyro/distributions/lkj.py#L67

```python
    def sample(self, sample_shape=torch.Size()):
        y = self._gen.sample(sample_shape=self.batch_shape + sample_shape).detach()
        z = y.mul(2).add(-1.0)
        return _vector_to_l_cholesky(z)
```

should be 

```python
    def sample(self, sample_shape=torch.Size()):
        y = self._gen.sample(sample_shape=sample_shape + self.batch_shape).detach()
        z = y.mul(2).add(-1.0)
        return _vector_to_l_cholesky(z)
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pyro/distributions/lkj.py`
Content:
```
1 # Copyright (c) 2017-2019 Uber Technologies, Inc.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 import math
5 
6 import torch
7 from torch.distributions import constraints
8 
9 from pyro.distributions.constraints import corr_cholesky_constraint
10 from pyro.distributions.torch import Beta
11 from pyro.distributions.torch_distribution import TorchDistribution
12 from pyro.distributions.transforms.cholesky import _vector_to_l_cholesky
13 
14 
15 # TODO: Modify class to support more than one eta value at a time?
16 class LKJCorrCholesky(TorchDistribution):
17     """
18     Generates cholesky factors of correlation matrices using an LKJ prior.
19 
20     The expected use is to combine it with a vector of variances and pass it
21     to the scale_tril parameter of a multivariate distribution such as MultivariateNormal.
22 
23     E.g., if theta is a (positive) vector of covariances with the same dimensionality
24     as this distribution, and Omega is sampled from this distribution,
25     scale_tril=torch.mm(torch.diag(sqrt(theta)), Omega)
26 
27     Note that the `event_shape` of this distribution is `[d, d]`
28 
29     .. note::
30 
31        When using this distribution with HMC/NUTS, it is important to
32        use a `step_size` such as 1e-4. If not, you are likely to experience LAPACK
33        errors regarding positive-definiteness.
34 
35     For example usage, refer to
36     `pyro/examples/lkj.py <https://github.com/pyro-ppl/pyro/blob/dev/examples/lkj.py>`_.
37 
38     :param int d: Dimensionality of the matrix
39     :param torch.Tensor eta: A single positive number parameterizing the distribution.
40     """
41     arg_constraints = {"eta": constraints.positive}
42     support = corr_cholesky_constraint
43     has_rsample = False
44 
45     def __init__(self, d, eta, validate_args=None):
46         if eta.numel() != 1:
47             raise ValueError("eta must be a single number; for a larger batch size, call expand")
48         if d <= 1:
49             raise ValueError("d must be > 1 in any correlation matrix")
50         eta = eta.squeeze()
51         vector_size = (d * (d - 1)) // 2
52         alpha = eta.add(0.5 * (d - 1.0))
53 
54         concentrations = torch.empty(vector_size, dtype=eta.dtype, device=eta.device)
55         i = 0
56         for k in range(d - 1):
57             alpha -= .5
58             concentrations[..., i:(i + d - k - 1)] = alpha
59             i += d - k - 1
60         self._gen = Beta(concentrations, concentrations)
61         self.eta = eta
62         self._d = d
63         self._lkj_constant = None
64         super().__init__(torch.Size(), torch.Size((d, d)), validate_args=validate_args)
65 
66     def sample(self, sample_shape=torch.Size()):
67         y = self._gen.sample(sample_shape=self.batch_shape + sample_shape).detach()
68         z = y.mul(2).add(-1.0)
69         return _vector_to_l_cholesky(z)
70 
71     def expand(self, batch_shape, _instance=None):
72         new = self._get_checked_instance(LKJCorrCholesky, _instance)
73         batch_shape = torch.Size(batch_shape)
74         new._gen = self._gen
75         new.eta = self.eta
76         new._d = self._d
77         new._lkj_constant = self._lkj_constant
78         super(LKJCorrCholesky, new).__init__(batch_shape, self.event_shape, validate_args=False)
79         new._validate_args = self._validate_args
80         return new
81 
82     def lkj_constant(self, eta, K):
83         if self._lkj_constant is not None:
84             return self._lkj_constant
85 
86         Km1 = K - 1
87 
88         constant = torch.lgamma(eta.add(0.5 * Km1)).mul(Km1)
89 
90         k = torch.linspace(start=1, end=Km1, steps=Km1, dtype=eta.dtype, device=eta.device)
91         constant -= (k.mul(math.log(math.pi) * 0.5) + torch.lgamma(eta.add(0.5 * (Km1 - k)))).sum()
92 
93         self._lkj_constant = constant
94         return constant
95 
96     def log_prob(self, x):
97         if self._validate_args:
98             self._validate_sample(x)
99 
100         eta = self.eta
101 
102         lp = self.lkj_constant(eta, self._d)
103 
104         Km1 = self._d - 1
105 
106         log_diagonals = x.diagonal(offset=0, dim1=-1, dim2=-2)[..., 1:].log()
107         # TODO: Figure out why the `device` kwarg to torch.linspace seems to not work in certain situations,
108         # and a seemingly redundant .to(x.device) is needed below.
109         values = log_diagonals * torch.linspace(start=Km1 - 1, end=0, steps=Km1,
110                                                 dtype=x.dtype,
111                                                 device=x.device).expand_as(log_diagonals).to(x.device)
112 
113         values += log_diagonals.mul(eta.mul(2).add(-2.0))
114         values = values.sum(-1) + lp
115         values, _ = torch.broadcast_tensors(values, torch.empty(self.batch_shape))
116         return values
117
```
--- 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/pyro/distributions/lkj.py b/pyro/distributions/lkj.py
--- a/pyro/distributions/lkj.py
+++ b/pyro/distributions/lkj.py
@@ -64,7 +64,8 @@
         super().__init__(torch.Size(), torch.Size((d, d)), validate_args=validate_args)
 
     def sample(self, sample_shape=torch.Size()):
-        y = self._gen.sample(sample_shape=self.batch_shape + sample_shape).detach()
+        with torch.no_grad():
+            y = self._gen.sample(sample_shape=sample_shape + self.batch_shape)
         z = y.mul(2).add(-1.0)
         return _vector_to_l_cholesky(z)

{"golden_diff": "diff --git a/pyro/distributions/lkj.py b/pyro/distributions/lkj.py\n--- a/pyro/distributions/lkj.py\n+++ b/pyro/distributions/lkj.py\n@@ -64,7 +64,8 @@\n         super().__init__(torch.Size(), torch.Size((d, d)), validate_args=validate_args)\n \n     def sample(self, sample_shape=torch.Size()):\n-        y = self._gen.sample(sample_shape=self.batch_shape + sample_shape).detach()\n+        with torch.no_grad():\n+            y = self._gen.sample(sample_shape=sample_shape + self.batch_shape)\n         z = y.mul(2).add(-1.0)\n         return _vector_to_l_cholesky(z)\n", "issue": "[bug] Samples from LKJCorrCholesky have the wrong shape\n### Issue Description\r\nThe batch and sample shapes for `LKJCorrCholesky` are swapped when sampling.\r\n\r\n### Environment\r\nOS: Ubuntu 16.04\r\nPython: 3.7.2\r\nTorch version: 1.6.0\r\nPyro version: 1.4.0\r\n\r\n### Code Snippet\r\nMinimal example:\r\n\r\n```python\r\n>>> dist = pyro.distributions.LKJCorrCholesky(d=3, eta=torch.ones(())).expand(12)\r\n>>> # batch shape and event shape are as you'd expect\r\n>>> dist.batch_shape\r\ntorch.Size([12])\r\n>>> dist.event_shape\r\ntorch.Size([3, 3])\r\n>>> # samples have correct shape when sample_shape=()\r\n>>> dist.shape(())\r\ntorch.Size([12, 3, 3])\r\n>>> dist.sample().shape\r\ntorch.Size([12, 3, 3])\r\n>>> # samples have the wrong shape when sample_shape is non-unit\r\n>>> dist.shape((4,))  # as expected\r\ntorch.Size([4, 12, 3, 3])\r\n>>> dist.sample((4,)).shape\r\ntorch.Size([12, 4, 3, 3])\r\n```\r\n\r\nI think this line is the culprit: https://github.com/pyro-ppl/pyro/blob/7c2c22c10dffda8a33ffbd593cc8d58819959e40/pyro/distributions/lkj.py#L67\r\n\r\n```python\r\n    def sample(self, sample_shape=torch.Size()):\r\n        y = self._gen.sample(sample_shape=self.batch_shape + sample_shape).detach()\r\n        z = y.mul(2).add(-1.0)\r\n        return _vector_to_l_cholesky(z)\r\n```\r\n\r\nshould be \r\n\r\n```python\r\n    def sample(self, sample_shape=torch.Size()):\r\n        y = self._gen.sample(sample_shape=sample_shape + self.batch_shape).detach()\r\n        z = y.mul(2).add(-1.0)\r\n        return _vector_to_l_cholesky(z)\r\n```\r\n\n", "before_files": [{"content": "# Copyright (c) 2017-2019 Uber Technologies, Inc.\n# SPDX-License-Identifier: Apache-2.0\n\nimport math\n\nimport torch\nfrom torch.distributions import constraints\n\nfrom pyro.distributions.constraints import corr_cholesky_constraint\nfrom pyro.distributions.torch import Beta\nfrom pyro.distributions.torch_distribution import TorchDistribution\nfrom pyro.distributions.transforms.cholesky import _vector_to_l_cholesky\n\n\n# TODO: Modify class to support more than one eta value at a time?\nclass LKJCorrCholesky(TorchDistribution):\n    \"\"\"\n    Generates cholesky factors of correlation matrices using an LKJ prior.\n\n    The expected use is to combine it with a vector of variances and pass it\n    to the scale_tril parameter of a multivariate distribution such as MultivariateNormal.\n\n    E.g., if theta is a (positive) vector of covariances with the same dimensionality\n    as this distribution, and Omega is sampled from this distribution,\n    scale_tril=torch.mm(torch.diag(sqrt(theta)), Omega)\n\n    Note that the `event_shape` of this distribution is `[d, d]`\n\n    .. note::\n\n       When using this distribution with HMC/NUTS, it is important to\n       use a `step_size` such as 1e-4. If not, you are likely to experience LAPACK\n       errors regarding positive-definiteness.\n\n    For example usage, refer to\n    `pyro/examples/lkj.py <https://github.com/pyro-ppl/pyro/blob/dev/examples/lkj.py>`_.\n\n    :param int d: Dimensionality of the matrix\n    :param torch.Tensor eta: A single positive number parameterizing the distribution.\n    \"\"\"\n    arg_constraints = {\"eta\": constraints.positive}\n    support = corr_cholesky_constraint\n    has_rsample = False\n\n    def __init__(self, d, eta, validate_args=None):\n        if eta.numel() != 1:\n            raise ValueError(\"eta must be a single number; for a larger batch size, call expand\")\n        if d <= 1:\n            raise ValueError(\"d must be > 1 in any correlation matrix\")\n        eta = eta.squeeze()\n        vector_size = (d * (d - 1)) // 2\n        alpha = eta.add(0.5 * (d - 1.0))\n\n        concentrations = torch.empty(vector_size, dtype=eta.dtype, device=eta.device)\n        i = 0\n        for k in range(d - 1):\n            alpha -= .5\n            concentrations[..., i:(i + d - k - 1)] = alpha\n            i += d - k - 1\n        self._gen = Beta(concentrations, concentrations)\n        self.eta = eta\n        self._d = d\n        self._lkj_constant = None\n        super().__init__(torch.Size(), torch.Size((d, d)), validate_args=validate_args)\n\n    def sample(self, sample_shape=torch.Size()):\n        y = self._gen.sample(sample_shape=self.batch_shape + sample_shape).detach()\n        z = y.mul(2).add(-1.0)\n        return _vector_to_l_cholesky(z)\n\n    def expand(self, batch_shape, _instance=None):\n        new = self._get_checked_instance(LKJCorrCholesky, _instance)\n        batch_shape = torch.Size(batch_shape)\n        new._gen = self._gen\n        new.eta = self.eta\n        new._d = self._d\n        new._lkj_constant = self._lkj_constant\n        super(LKJCorrCholesky, new).__init__(batch_shape, self.event_shape, validate_args=False)\n        new._validate_args = self._validate_args\n        return new\n\n    def lkj_constant(self, eta, K):\n        if self._lkj_constant is not None:\n            return self._lkj_constant\n\n        Km1 = K - 1\n\n        constant = torch.lgamma(eta.add(0.5 * Km1)).mul(Km1)\n\n        k = torch.linspace(start=1, end=Km1, steps=Km1, dtype=eta.dtype, device=eta.device)\n        constant -= (k.mul(math.log(math.pi) * 0.5) + torch.lgamma(eta.add(0.5 * (Km1 - k)))).sum()\n\n        self._lkj_constant = constant\n        return constant\n\n    def log_prob(self, x):\n        if self._validate_args:\n            self._validate_sample(x)\n\n        eta = self.eta\n\n        lp = self.lkj_constant(eta, self._d)\n\n        Km1 = self._d - 1\n\n        log_diagonals = x.diagonal(offset=0, dim1=-1, dim2=-2)[..., 1:].log()\n        # TODO: Figure out why the `device` kwarg to torch.linspace seems to not work in certain situations,\n        # and a seemingly redundant .to(x.device) is needed below.\n        values = log_diagonals * torch.linspace(start=Km1 - 1, end=0, steps=Km1,\n                                                dtype=x.dtype,\n                                                device=x.device).expand_as(log_diagonals).to(x.device)\n\n        values += log_diagonals.mul(eta.mul(2).add(-2.0))\n        values = values.sum(-1) + lp\n        values, _ = torch.broadcast_tensors(values, torch.empty(self.batch_shape))\n        return values\n", "path": "pyro/distributions/lkj.py"}], "after_files": [{"content": "# Copyright (c) 2017-2019 Uber Technologies, Inc.\n# SPDX-License-Identifier: Apache-2.0\n\nimport math\n\nimport torch\nfrom torch.distributions import constraints\n\nfrom pyro.distributions.constraints import corr_cholesky_constraint\nfrom pyro.distributions.torch import Beta\nfrom pyro.distributions.torch_distribution import TorchDistribution\nfrom pyro.distributions.transforms.cholesky import _vector_to_l_cholesky\n\n\n# TODO: Modify class to support more than one eta value at a time?\nclass LKJCorrCholesky(TorchDistribution):\n    \"\"\"\n    Generates cholesky factors of correlation matrices using an LKJ prior.\n\n    The expected use is to combine it with a vector of variances and pass it\n    to the scale_tril parameter of a multivariate distribution such as MultivariateNormal.\n\n    E.g., if theta is a (positive) vector of covariances with the same dimensionality\n    as this distribution, and Omega is sampled from this distribution,\n    scale_tril=torch.mm(torch.diag(sqrt(theta)), Omega)\n\n    Note that the `event_shape` of this distribution is `[d, d]`\n\n    .. note::\n\n       When using this distribution with HMC/NUTS, it is important to\n       use a `step_size` such as 1e-4. If not, you are likely to experience LAPACK\n       errors regarding positive-definiteness.\n\n    For example usage, refer to\n    `pyro/examples/lkj.py <https://github.com/pyro-ppl/pyro/blob/dev/examples/lkj.py>`_.\n\n    :param int d: Dimensionality of the matrix\n    :param torch.Tensor eta: A single positive number parameterizing the distribution.\n    \"\"\"\n    arg_constraints = {\"eta\": constraints.positive}\n    support = corr_cholesky_constraint\n    has_rsample = False\n\n    def __init__(self, d, eta, validate_args=None):\n        if eta.numel() != 1:\n            raise ValueError(\"eta must be a single number; for a larger batch size, call expand\")\n        if d <= 1:\n            raise ValueError(\"d must be > 1 in any correlation matrix\")\n        eta = eta.squeeze()\n        vector_size = (d * (d - 1)) // 2\n        alpha = eta.add(0.5 * (d - 1.0))\n\n        concentrations = torch.empty(vector_size, dtype=eta.dtype, device=eta.device)\n        i = 0\n        for k in range(d - 1):\n            alpha -= .5\n            concentrations[..., i:(i + d - k - 1)] = alpha\n            i += d - k - 1\n        self._gen = Beta(concentrations, concentrations)\n        self.eta = eta\n        self._d = d\n        self._lkj_constant = None\n        super().__init__(torch.Size(), torch.Size((d, d)), validate_args=validate_args)\n\n    def sample(self, sample_shape=torch.Size()):\n        with torch.no_grad():\n            y = self._gen.sample(sample_shape=sample_shape + self.batch_shape)\n        z = y.mul(2).add(-1.0)\n        return _vector_to_l_cholesky(z)\n\n    def expand(self, batch_shape, _instance=None):\n        new = self._get_checked_instance(LKJCorrCholesky, _instance)\n        batch_shape = torch.Size(batch_shape)\n        new._gen = self._gen\n        new.eta = self.eta\n        new._d = self._d\n        new._lkj_constant = self._lkj_constant\n        super(LKJCorrCholesky, new).__init__(batch_shape, self.event_shape, validate_args=False)\n        new._validate_args = self._validate_args\n        return new\n\n    def lkj_constant(self, eta, K):\n        if self._lkj_constant is not None:\n            return self._lkj_constant\n\n        Km1 = K - 1\n\n        constant = torch.lgamma(eta.add(0.5 * Km1)).mul(Km1)\n\n        k = torch.linspace(start=1, end=Km1, steps=Km1, dtype=eta.dtype, device=eta.device)\n        constant -= (k.mul(math.log(math.pi) * 0.5) + torch.lgamma(eta.add(0.5 * (Km1 - k)))).sum()\n\n        self._lkj_constant = constant\n        return constant\n\n    def log_prob(self, x):\n        if self._validate_args:\n            self._validate_sample(x)\n\n        eta = self.eta\n\n        lp = self.lkj_constant(eta, self._d)\n\n        Km1 = self._d - 1\n\n        log_diagonals = x.diagonal(offset=0, dim1=-1, dim2=-2)[..., 1:].log()\n        # TODO: Figure out why the `device` kwarg to torch.linspace seems to not work in certain situations,\n        # and a seemingly redundant .to(x.device) is needed below.\n        values = log_diagonals * torch.linspace(start=Km1 - 1, end=0, steps=Km1,\n                                                dtype=x.dtype,\n                                                device=x.device).expand_as(log_diagonals).to(x.device)\n\n        values += log_diagonals.mul(eta.mul(2).add(-2.0))\n        values = values.sum(-1) + lp\n        values, _ = torch.broadcast_tensors(values, torch.empty(self.batch_shape))\n        return values\n", "path": "pyro/distributions/lkj.py"}]}