Make TransformReparam compatible with .to_event()

pyro/infer/reparam/transform.py

@@ -20,15 +20,23 @@ class TransformReparam(Reparam):
     """
     def __call__(self, name, fn, obs):
         assert obs is None, "TransformReparam does not support observe statements"
+        fn, event_dim = self._unwrap(fn)
         assert isinstance(fn, dist.TransformedDistribution)
 
         # Draw noise from the base distribution.
-        x = pyro.sample("{}_base".format(name), fn.base_dist)
+        base_event_dim = event_dim
+        try:  # requires https://github.com/pyro-ppl/pyro/pull/2739
+            for t in reversed(fn.transforms):
+                base_event_dim += t.domain.event_dim - t.codomain.event_dim
+        except AttributeError:
+            pass
+        x = pyro.sample("{}_base".format(name),
+                        self._wrap(fn.base_dist, base_event_dim))
 
         # Differentiably transform.
         for t in fn.transforms:
             x = t(x)
 
         # Simulate a pyro.deterministic() site.
-        new_fn = dist.Delta(x, event_dim=fn.event_dim)
+        new_fn = dist.Delta(x, event_dim=event_dim).mask(False)
         return new_fn, x

tests/infer/reparam/test_transform.py

@@ -27,29 +27,32 @@ def get_moments(x):
     return torch.stack([m1, m2, m3, m4])
 
 
-@pytest.mark.parametrize("shape", [(), (4,), (2, 3)], ids=str)
-def test_log_normal(shape):
+@pytest.mark.parametrize("batch_shape", [(), (4,), (2, 3)], ids=str)
+@pytest.mark.parametrize("event_shape", [(), (5,)], ids=str)
+def test_log_normal(batch_shape, event_shape):
+    shape = batch_shape + event_shape
     loc = torch.empty(shape).uniform_(-1, 1)
     scale = torch.empty(shape).uniform_(0.5, 1.5)
 
     def model():
-        with pyro.plate_stack("plates", shape):
+        fn = dist.TransformedDistribution(
+            dist.Normal(torch.zeros_like(loc), torch.ones_like(scale)),
+            [AffineTransform(loc, scale), ExpTransform()])
+        if event_shape:
+            fn = fn.to_event(len(event_shape))
+        with pyro.plate_stack("plates", batch_shape):
             with pyro.plate("particles", 200000):
-                return pyro.sample("x",
-                                   dist.TransformedDistribution(
-                                       dist.Normal(torch.zeros_like(loc),
-                                                   torch.ones_like(scale)),
-                                       [AffineTransform(loc, scale),
-                                        ExpTransform()]))
+                return pyro.sample("x", fn)
 
     with poutine.trace() as tr:
         value = model()
-    assert isinstance(tr.trace.nodes["x"]["fn"], dist.TransformedDistribution)
+    assert isinstance(tr.trace.nodes["x"]["fn"],
+                      (dist.TransformedDistribution, dist.Independent))
     expected_moments = get_moments(value)
 
     with poutine.reparam(config={"x": TransformReparam()}):
         with poutine.trace() as tr:
             value = model()
-    assert isinstance(tr.trace.nodes["x"]["fn"], dist.Delta)
+    assert isinstance(tr.trace.nodes["x"]["fn"], (dist.Delta, dist.MaskedDistribution))
     actual_moments = get_moments(value)
     assert_close(actual_moments, expected_moments, atol=0.05)