Conditional flows

README.md

@@ -32,6 +32,19 @@ base_distribution = distributions.StandardNormal(shape=[2])
 flow = flows.Flow(transform=transform, distribution=base_distribution)
 ```
 
+or use already implemented flow architectures:
+
+```python
+from nflows.flows import MaskedAutoregressiveFlow, SimpleRealNVP
+
+features=2
+hidden_features=4
+
+maf = MaskedAutoregressiveFlow(features=features, hidden_features=hidden_features)
+rnvp = SimpleRealNVP(features=features, hidden_features=hidden_features)
+nice = SimpleRealNVP(feautres=features, hidden_features=hidden_features, use_volumne_perserving=True)
+```
+
 To evaluate log probabilities of inputs:
 ```python
 log_prob = flow.log_prob(inputs)

environment.yml

@@ -28,4 +28,3 @@ dependencies:
   - pyyaml
   - tensorboard
   - tqdm
-

nflows/flows/autoregressive.py

@@ -1,6 +1,6 @@
 """Implementations of autoregressive flows."""
 
-from torch.nn import functional as F
+import torch.nn.functional as F
 
 from nflows.distributions.normal import StandardNormal
 from nflows.flows.base import Flow
@@ -19,18 +19,19 @@ class MaskedAutoregressiveFlow(Flow):
     """
 
     def __init__(
-        self,
-        features,
-        hidden_features,
-        num_layers,
-        num_blocks_per_layer,
-        use_residual_blocks=True,
-        use_random_masks=False,
-        use_random_permutations=False,
-        activation=F.relu,
-        dropout_probability=0.0,
-        batch_norm_within_layers=False,
-        batch_norm_between_layers=False,
+            self,
+            features,
+            hidden_features,
+            context_features=None,
+            num_layers=5,
+            num_blocks_per_layer=2,
+            use_residual_blocks=True,
+            use_random_masks=False,
+            use_random_permutations=False,
+            activation=F.relu,
+            dropout_probability=0.0,
+            batch_norm_within_layers=False,
+            batch_norm_between_layers=False,
     ):
 
         if use_random_permutations:
@@ -45,6 +46,7 @@ def __init__(
                 MaskedAffineAutoregressiveTransform(
                     features=features,
                     hidden_features=hidden_features,
+                    context_features=context_features,
                     num_blocks=num_blocks_per_layer,
                     use_residual_blocks=use_residual_blocks,
                     random_mask=use_random_masks,

nflows/flows/realnvp.py

@@ -24,16 +24,17 @@ class SimpleRealNVP(Flow):
     """
 
     def __init__(
-        self,
-        features,
-        hidden_features,
-        num_layers,
-        num_blocks_per_layer,
-        use_volume_preserving=False,
-        activation=F.relu,
-        dropout_probability=0.0,
-        batch_norm_within_layers=False,
-        batch_norm_between_layers=False,
+            self,
+            features,
+            hidden_features,
+            context_features=None,
+            num_layers=5,
+            num_blocks_per_layer=2,
+            use_volume_preserving=False,
+            activation=F.relu,
+            dropout_probability=0.0,
+            batch_norm_within_layers=False,
+            batch_norm_between_layers=False,
     ):
 
         if use_volume_preserving:
@@ -49,6 +50,7 @@ def create_resnet(in_features, out_features):
                 in_features,
                 out_features,
                 hidden_features=hidden_features,
+                context_features=context_features,
                 num_blocks=num_blocks_per_layer,
                 activation=activation,
                 dropout_probability=dropout_probability,

tests/flows/autoregressive_test.py

@@ -30,6 +30,33 @@ def test_sample(self):
         self.assertIsInstance(samples, torch.Tensor)
         self.assertEqual(samples.shape, torch.Size([num_samples, features]))
 
+    def test_conditional_log_prob(self):
+        batch_size = 10
+        features = 20
+        context_features = 5
+        flow = ar.MaskedAutoregressiveFlow(
+            features=features, hidden_features=30, context_features=context_features
+        )
+        inputs = torch.randn(batch_size, features)
+        context = torch.randn(batch_size, context_features)
+        log_prob = flow.log_prob(inputs, context)
+        self.assertIsInstance(log_prob, torch.Tensor)
+        self.assertEqual(log_prob.shape, torch.Size([batch_size]))
+
+    def test_conditional_sample(self):
+        num_samples = 10
+
+        batch_size = 10
+        features = 20
+        context_features = 5
+        flow = ar.MaskedAutoregressiveFlow(
+            features=features, hidden_features=30, context_features=context_features
+        )
+        context = torch.randn(batch_size, context_features)
+        samples = flow.sample(num_samples, context)
+        self.assertIsInstance(samples, torch.Tensor)
+        self.assertEqual(samples.shape, torch.Size([batch_size, num_samples, features]))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/flows/realnvp_test.py

@@ -30,6 +30,33 @@ def test_sample(self):
         self.assertIsInstance(samples, torch.Tensor)
         self.assertEqual(samples.shape, torch.Size([num_samples, features]))
 
+    def test_conditional_log_prob(self):
+        batch_size = 10
+        features = 20
+        context_features = 5
+        flow = realnvp.SimpleRealNVP(
+            features=features, hidden_features=30, context_features=context_features
+        )
+        inputs = torch.randn(batch_size, features)
+        context = torch.randn(batch_size, context_features)
+        log_prob = flow.log_prob(inputs, context)
+        self.assertIsInstance(log_prob, torch.Tensor)
+        self.assertEqual(log_prob.shape, torch.Size([batch_size]))
+
+    def test_conditional_sample(self):
+        num_samples = 10
+
+        batch_size = 10
+        features = 20
+        context_features = 5
+        flow = realnvp.SimpleRealNVP(
+            features=features, hidden_features=30, context_features=context_features
+        )
+        context = torch.randn(batch_size, context_features)
+        samples = flow.sample(num_samples, context)
+        self.assertIsInstance(samples, torch.Tensor)
+        self.assertEqual(samples.shape, torch.Size([batch_size, num_samples, features]))
+
 
 if __name__ == "__main__":
     unittest.main()