Fix shape error when sampling

normalizing_flows/flows.py

@@ -5,7 +5,7 @@
 from normalizing_flows.bijections.base import Bijection
 from normalizing_flows.bijections.continuous.ddnf import DeepDiffeomorphicBijection
 from normalizing_flows.regularization import reconstruction_error
-from normalizing_flows.utils import flatten_event, get_batch_shape
+from normalizing_flows.utils import flatten_event, get_batch_shape, unflatten_event
 
 
 class Flow(nn.Module):
@@ -24,6 +24,13 @@ def base_log_prob(self, z):
         log_prob = self.base.log_prob(zf)
         return log_prob
 
+    def base_sample(self, sample_shape):
+        z_flat = self.base.sample(sample_shape)
+        z = unflatten_event(z_flat, self.bijection.event_shape)
+        return z
+
+
+
     def forward_with_log_prob(self, x: torch.Tensor, context: torch.Tensor = None):
         if context is not None:
             assert context.shape[0] == x.shape[0]
@@ -46,11 +53,11 @@ def sample(self, n: int, context: torch.Tensor = None, no_grad: bool = False, re
         :return:
         """
         if context is not None:
-            z = self.base.sample(sample_shape=torch.Size((n, len(context))))
+            z = self.base_sample(sample_shape=torch.Size((n, len(context))))
             context = context[None].repeat(*[n, *([1] * len(context.shape))])  # Make context shape match z shape
             assert z.shape[:2] == context.shape[:2]
         else:
-            z = self.base.sample(sample_shape=torch.Size((n,)))
+            z = self.base_sample(sample_shape=torch.Size((n,)))
         if no_grad:
             z = z.detach()
             with torch.no_grad():

test/test_sample.py

@@ -0,0 +1,11 @@
+import torch
+
+from normalizing_flows import Flow
+from normalizing_flows.bijections import RealNVP
+
+
+def test_real_nvp():
+    torch.manual_seed(0)
+    f = Flow(RealNVP(event_shape=torch.Size((2, 3, 5, 7))))
+    x = f.sample(10)
+    assert x.shape == (10, 2, 3, 5, 7)