Repository: pyro-ppl/pyro

Issue:
guide.quantiles() shape issue
I run into frequent `RuntimeError`s when generating quantiles of models fitted with SVI. This can be replicated on the radon models in `pyro-models`. For example, the [radon_intercept](https://github.com/pyro-ppl/pyro-models/blob/master/pyro_models/arm/radon_intercept_chr.py) model fit with `SVI` with an `AutoNormal` guide generates the following:

```
>>> estimates = guide.quantiles([0.05, 0.5, 0.95])


estimates = guide.quantiles([0.05, 0.5, 0.95])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
~/GitHub/pyro-models/pyro_models/arm/radon_intercept_chr.py in 
----> 78 estimates = guide.quantiles([0.05, 0.5, 0.95])

~/miniforge3/envs/pyro/lib/python3.9/site-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     25         def decorate_context(*args, **kwargs):
     26             with self.__class__():
---> 27                 return func(*args, **kwargs)
     28         return cast(F, decorate_context)
     29 

~/miniforge3/envs/pyro/lib/python3.9/site-packages/pyro/infer/autoguide/guides.py in quantiles(self, quantiles, *args, **kwargs)
    578 
    579             site_quantiles = torch.tensor(quantiles, dtype=site_loc.dtype, device=site_loc.device)
--> 580             site_quantiles_values = dist.Normal(site_loc, site_scale).icdf(site_quantiles)
    581             constrained_site_quantiles = biject_to(site["fn"].support)(site_quantiles_values)
    582             results[name] = constrained_site_quantiles

~/miniforge3/envs/pyro/lib/python3.9/site-packages/torch/distributions/normal.py in icdf(self, value)
     83 
     84     def icdf(self, value):
---> 85         return self.loc + self.scale * torch.erfinv(2 * value - 1) * math.sqrt(2)
     86 
     87     def entropy(self):

RuntimeError: The size of tensor a (85) must match the size of tensor b (3) at non-singleton dimension 0

```

This has come up repeatedly with larger models, but this is the first easily reproducible example that I have come across.

### Environment
For any bugs, please provide the following:
 - OS and python version. MacOS, Linux; Python 3.7.10
 - PyTorch version: 1.8.0
 - Pyro version: 1.6.0

### Code Snippet

Using the `arm.radon_intercept` model, I am fitting using the following:

```python
N_STEPS = 1000

guide = autoguides.AutoNormal(model)

initial_lr = 0.01
optimizer = pyro.optim.ClippedAdam({"lr": initial_lr})
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

pyro.clear_param_store()
for j in range(N_STEPS):
    # calculate the loss and take a gradient step
    loss = svi.step(radon_data)
    if not j % 100:
        print("[iteration %04d] loss: %.4f" % (j + 1, loss / len(radon_data['y'])))

estimates = guide.quantiles([0.05, 0.5, 0.95])
```


Discussion:
Hi @fonnesbeck it looks like this is a bug in the model. To fix locally I believe you can fix the indexing, something like this:
```diff
  def model(data, params):
      # initialize data
      J = data["J"]
      N = data["N"]
      county = data["county"].long() - 1
      y = data["y"]
      sigma_eta =  pyro.sample("sigma_eta", dist.Uniform(0., 100.))
      sigma_y =  pyro.sample("sigma_y", dist.Uniform(0., 100.))
      mu =  pyro.sample("mu", dist.Normal(0., 1.))
      with pyro.plate("J", J):
          eta =  pyro.sample("eta", dist.Normal(0., sigma_eta))
      with pyro.plate("data", N):
-         y_hat = 0.1 * mu + eta[county]
+         y_hat = 0.1 * mu + eta[..., county]  # or something like this
          y =  pyro.sample("y", dist.Normal(y_hat, sigma_y), obs=y)
```
The issues is that during quantiles (or any other multiple-sample method), `eta` is batched, and the naive `eta[county]` indexes into the wrong dimension.
that's what i thought at first too but i don't think that's actually it since `quantiles` doesn't actually run the `model` afaik...
My guess was that it was something deeper, but its just a guess. I've been building Pyro models for the past 6 weeks or so and I can only recall one occasion where `quantiles` worked without this error. Of course, I may be constructing all of my models incorrectly ...
I think it is a bug. The line
```
--> 580             site_quantiles_values = dist.Normal(site_loc, site_scale).icdf(site_quantiles)
```
assumes that quantiles is a scalar. I'll add a fix for it.

Resolving patch:
diff --git a/pyro/infer/autoguide/guides.py b/pyro/infer/autoguide/guides.py
--- a/pyro/infer/autoguide/guides.py
+++ b/pyro/infer/autoguide/guides.py
@@ -565,7 +565,7 @@ def quantiles(self, quantiles, *args, **kwargs):
 
         :param quantiles: A list of requested quantiles between 0 and 1.
         :type quantiles: torch.Tensor or list
-        :return: A dict mapping sample site name to a list of quantile values.
+        :return: A dict mapping sample site name to a tensor of quantile values.
         :rtype: dict
         """
         results = {}
@@ -574,6 +574,7 @@ def quantiles(self, quantiles, *args, **kwargs):
             site_loc, site_scale = self._get_loc_and_scale(name)
 
             site_quantiles = torch.tensor(quantiles, dtype=site_loc.dtype, device=site_loc.device)
+            site_quantiles = site_quantiles.reshape((-1,) + (1,) * site_loc.dim())
             site_quantiles_values = dist.Normal(site_loc, site_scale).icdf(site_quantiles)
             constrained_site_quantiles = biject_to(site["fn"].support)(site_quantiles_values)
             results[name] = constrained_site_quantiles
@@ -782,7 +783,7 @@ def quantiles(self, quantiles, *args, **kwargs):
 
         :param quantiles: A list of requested quantiles between 0 and 1.
         :type quantiles: torch.Tensor or list
-        :return: A dict mapping sample site name to a list of quantile values.
+        :return: A dict mapping sample site name to a tensor of quantile values.
         :rtype: dict
         """
         loc, scale = self._loc_scale(*args, **kwargs)
@@ -792,6 +793,7 @@ def quantiles(self, quantiles, *args, **kwargs):
         for latent in latents:
             for site, unconstrained_value in self._unpack_latent(latent):
                 result.setdefault(site["name"], []).append(biject_to(site["fn"].support)(unconstrained_value))
+        result = {k: torch.stack(v) for k, v in result.items()}
         return result
 
 


Test patch:
diff --git a/tests/infer/test_autoguide.py b/tests/infer/test_autoguide.py
--- a/tests/infer/test_autoguide.py
+++ b/tests/infer/test_autoguide.py
@@ -404,7 +404,7 @@ def test_quantiles(auto_class, Elbo):
     def model():
         pyro.sample("x", dist.Normal(0.0, 1.0))
         pyro.sample("y", dist.LogNormal(0.0, 1.0))
-        pyro.sample("z", dist.Beta(2.0, 2.0))
+        pyro.sample("z", dist.Beta(2.0, 2.0).expand([2]).to_event(1))
 
     guide = auto_class(model)
     optim = Adam({'lr': 0.05, 'betas': (0.8, 0.99)})
@@ -421,10 +421,12 @@ def model():
     median = guide.median()
     for name in ["x", "y", "z"]:
         assert not median[name].requires_grad
+        assert torch.is_tensor(quantiles[name])
+        assert len(quantiles[name]) == 3
         for q in quantiles[name]:
             assert not q.requires_grad
+            assert q.shape == (() if name != "z" else (2,))
         assert_equal(median[name], quantiles[name][1])
-    quantiles = {name: [v.item() for v in value] for name, value in quantiles.items()}
 
     assert -3.0 < quantiles["x"][0]
     assert quantiles["x"][0] + 1.0 < quantiles["x"][1]
@@ -436,10 +438,10 @@ def model():
     assert quantiles["y"][1] * 2.0 < quantiles["y"][2]
     assert quantiles["y"][2] < 100.0
 
-    assert 0.01 < quantiles["z"][0]
-    assert quantiles["z"][0] + 0.1 < quantiles["z"][1]
-    assert quantiles["z"][1] + 0.1 < quantiles["z"][2]
-    assert quantiles["z"][2] < 0.99
+    assert (0.01 < quantiles["z"][0]).all()
+    assert (quantiles["z"][0] + 0.1 < quantiles["z"][1]).all()
+    assert (quantiles["z"][1] + 0.1 < quantiles["z"][2]).all()
+    assert (quantiles["z"][2] < 0.99).all()
 
 
 @pytest.mark.parametrize("continuous_class", [