Fix autoguide latent shape mismatch (#1961)

pyro/contrib/autoguide/__init__.py

@@ -403,10 +403,12 @@ def _unpack_latent(self, latent):
         batch_shape = latent.shape[:-1]  # for plates outside of _setup_prototype, e.g. parallel particles
         pos = 0
         for name, site in self.prototype_trace.iter_stochastic_nodes():
+            constrained_shape = site["value"].shape
             unconstrained_shape = self._unconstrained_shapes[name]
             size = _product(unconstrained_shape)
+            event_dim = site["fn"].event_dim + len(unconstrained_shape) - len(constrained_shape)
             unconstrained_shape = broadcast_shape(unconstrained_shape,
-                                                  batch_shape + (1,) * site["fn"].event_dim)
+                                                  batch_shape + (1,) * event_dim)
             unconstrained_value = latent[..., pos:pos + size].view(unconstrained_shape)
             yield site, unconstrained_value
             pos += size

tests/contrib/autoguide/test_advi.py

@@ -314,3 +314,13 @@ def model():
     guide = AutoDiagonalNormal(model)
     with pytest.raises(RuntimeError):
         guide()
+
+
+def test_unpack_latent():
+    def model():
+        return pyro.sample('x', dist.LKJCorrCholesky(2, torch.tensor(1.)))
+
+    guide = AutoDiagonalNormal(model)
+    assert guide()['x'].shape == model().shape
+    latent = guide.sample_latent()
+    assert list(guide._unpack_latent(latent))[0][1].shape == (1,)