Old test module replaced

normflows/flows/affine/autoregressive_test.py

@@ -2,12 +2,12 @@
 import unittest
 
 from normflows.flows.affine import autoregressive
-from normflows.flows.neural_spline.flow_test import FlowTest
+from normflows.flows.flow_test import FlowTest
 
 
 class MaskedAffineAutoregressiveTest(FlowTest):
-    def test_forward(self):
-        batch_size = 10
+    def test_maf(self):
+        batch_size = 3
         features = 20
         inputs = torch.randn(batch_size, features)
         for use_residual_blocks, random_mask in [
@@ -18,61 +18,14 @@ def test_forward(self):
             with self.subTest(
                 use_residual_blocks=use_residual_blocks, random_mask=random_mask
             ):
-                transform = autoregressive.MaskedAffineAutoregressive(
+                flow = autoregressive.MaskedAffineAutoregressive(
                     features=features,
                     hidden_features=30,
                     num_blocks=5,
                     use_residual_blocks=use_residual_blocks,
                     random_mask=random_mask,
                 )
-                outputs, logabsdet = transform(inputs)
-                self.assert_tensor_is_good(outputs, [batch_size, features])
-                self.assert_tensor_is_good(logabsdet, [batch_size])
-
-    def test_inverse(self):
-        batch_size = 10
-        features = 20
-        inputs = torch.randn(batch_size, features)
-        for use_residual_blocks, random_mask in [
-            (False, False),
-            (False, True),
-            (True, False),
-        ]:
-            with self.subTest(
-                use_residual_blocks=use_residual_blocks, random_mask=random_mask
-            ):
-                transform = autoregressive.MaskedAffineAutoregressive(
-                    features=features,
-                    hidden_features=30,
-                    num_blocks=5,
-                    use_residual_blocks=use_residual_blocks,
-                    random_mask=random_mask,
-                )
-                outputs, logabsdet = transform.inverse(inputs)
-                self.assert_tensor_is_good(outputs, [batch_size, features])
-                self.assert_tensor_is_good(logabsdet, [batch_size])
-
-    def test_forward_inverse_are_consistent(self):
-        batch_size = 10
-        features = 20
-        inputs = torch.randn(batch_size, features)
-        self.eps = 1e-6
-        for use_residual_blocks, random_mask in [
-            (False, False),
-            (False, True),
-            (True, False),
-        ]:
-            with self.subTest(
-                use_residual_blocks=use_residual_blocks, random_mask=random_mask
-            ):
-                transform = autoregressive.MaskedAffineAutoregressive(
-                    features=features,
-                    hidden_features=30,
-                    num_blocks=5,
-                    use_residual_blocks=use_residual_blocks,
-                    random_mask=random_mask,
-                )
-                self.assert_forward_inverse_are_consistent(transform, inputs)
+                self.checkForwardInverse(flow, inputs)
 
 
 if __name__ == "__main__":

normflows/flows/affine/glow_test.py

@@ -22,11 +22,11 @@ def test_glow(self):
             with self.subTest(batch_size=batch_size, channels=channels,
                               scale=scale, split_mode=split_mode,
                               use_lu=use_lu, net_actnorm=net_actnorm):
-                input = torch.rand((batch_size, channels) + img_size)
+                inputs = torch.rand((batch_size, channels) + img_size)
                 flow = GlowBlock(channels, hidden_channels,
                                  scale=scale, split_mode=split_mode,
                                  use_lu=use_lu, net_actnorm=net_actnorm)
-                self.checkForwardInverse(flow, input)
+                self.checkForwardInverse(flow, inputs)
 
 
 if __name__ == "__main__":

normflows/flows/flow_test.py

@@ -11,32 +11,32 @@ class FlowTest(unittest.TestCase):
     def assertClose(self, actual, expected):
         assert_close(actual, expected)
 
-    def checkForward(self, flow, input):
+    def checkForward(self, flow, inputs):
         # Do forward transform
-        output, log_det = flow(input)
+        outputs, log_det = flow(inputs)
         # Check type
-        assert output.dtype == input.dtype
+        assert outputs.dtype == inputs.dtype
         # Check shape
-        assert output.shape == input.shape
+        assert outputs.shape == inputs.shape
         # Return results
-        return output, log_det
+        return outputs, log_det
 
-    def checkInverse(self, flow, input):
+    def checkInverse(self, flow, inputs):
         # Do inverse transform
-        output, log_det = flow.inverse(input)
+        outputs, log_det = flow.inverse(inputs)
         # Check type
-        assert output.dtype == input.dtype
+        assert outputs.dtype == inputs.dtype
         # Check shape
-        assert output.shape == input.shape
+        assert outputs.shape == inputs.shape
         # Return results
-        return output, log_det
+        return outputs, log_det
 
-    def checkForwardInverse(self, flow, input):
+    def checkForwardInverse(self, flow, inputs):
         # Check forward
-        output, log_det = self.checkForward(flow, input)
+        outputs, log_det = self.checkForward(flow, inputs)
         # Check inverse
-        input_, log_det_ = self.checkInverse(flow, output)
+        input_, log_det_ = self.checkInverse(flow, outputs)
         # Check identity
-        self.assertClose(input_, input)
+        self.assertClose(input_, inputs)
         ld_id = log_det + log_det_
         self.assertClose(ld_id, torch.zeros_like(ld_id))

normflows/flows/neural_spline/autoregressive_test.py

@@ -1,31 +1,30 @@
 """
 Tests for the autoregressive transforms.
-Code taken from https://github.com/bayesiains/nsf
+Code partially taken from https://github.com/bayesiains/nsf
 """
 
 import torch
 import unittest
 
 from normflows.flows.neural_spline import autoregressive
-from normflows.flows.neural_spline.flow_test import FlowTest
+from normflows.flows.flow_test import FlowTest
 
 
 class MaskedPiecewiseRationalQuadraticAutoregressiveFlowTest(FlowTest):
-    def test_forward_inverse_are_consistent(self):
-        batch_size = 10
-        features = 20
+    def test_mprqas(self):
+        batch_size = 5
+        features = 10
         inputs = torch.rand(batch_size, features)
-        self.eps = 1e-3
 
-        transform = autoregressive.MaskedPiecewiseRationalQuadraticAutoregressive(
+        flow = autoregressive.MaskedPiecewiseRationalQuadraticAutoregressive(
             num_bins=10,
             features=features,
             hidden_features=30,
             num_blocks=5,
             use_residual_blocks=True,
         )
 
-        self.assert_forward_inverse_are_consistent(transform, inputs)
+        self.checkForwardInverse(flow, inputs)
 
 
 if __name__ == "__main__":

normflows/flows/neural_spline/coupling_test.py

@@ -1,6 +1,6 @@
 """
 Tests for the coupling Transforms.
-Code taken from https://github.com/bayesiains/nsf
+Code partially taken from https://github.com/bayesiains/nsf
 """
 
 import torch
@@ -9,11 +9,11 @@
 from normflows import nets as nn_
 
 from normflows.flows.neural_spline import coupling
-from normflows.flows.neural_spline.flow_test import FlowTest
+from normflows.flows.flow_test import FlowTest
 from normflows import utils
 
 
-def create_coupling_transform(cls, shape, **kwargs):
+def create_coupling_transform(shape, **kwargs):
     if len(shape) == 1:
 
         def create_net(in_features, out_features):
@@ -31,104 +31,32 @@ def create_net(in_channels, out_channels):
 
     mask = utils.masks.create_mid_split_binary_mask(shape[0])
 
-    return cls(mask=mask, transform_net_create_fn=create_net, **kwargs), mask
+    return coupling.PiecewiseRationalQuadraticCoupling(mask=mask, transform_net_create_fn=create_net, **kwargs), mask
 
 
-batch_size = 10
+batch_size = 5
 
 
 class PiecewiseCouplingTransformTest(FlowTest):
-    classes = [coupling.PiecewiseRationalQuadraticCoupling]
-
     shapes = [[20], [2, 4, 4]]
 
-    def test_forward(self):
-        for shape in self.shapes:
-            for cls in self.classes:
-                inputs = torch.rand(batch_size, *shape)
-                transform, mask = create_coupling_transform(cls, shape)
-                outputs, logabsdet = transform(inputs)
-                with self.subTest(cls=cls, shape=shape):
-                    self.assert_tensor_is_good(outputs, [batch_size] + shape)
-                    self.assert_tensor_is_good(logabsdet, [batch_size])
-                    self.assertEqual(
-                        outputs[:, mask <= 0, ...], inputs[:, mask <= 0, ...]
-                    )
-
-    def test_forward_unconstrained(self):
-        batch_size = 10
+    def test_rqcs(self):
         for shape in self.shapes:
-            for cls in self.classes:
-                inputs = 3.0 * torch.randn(batch_size, *shape)
-                transform, mask = create_coupling_transform(cls, shape, tails="linear")
-                outputs, logabsdet = transform(inputs)
-                with self.subTest(cls=cls, shape=shape):
-                    self.assert_tensor_is_good(outputs, [batch_size] + shape)
-                    self.assert_tensor_is_good(logabsdet, [batch_size])
-                    self.assertEqual(
-                        outputs[:, mask <= 0, ...], inputs[:, mask <= 0, ...]
-                    )
-
-    def test_inverse(self):
-        for shape in self.shapes:
-            for cls in self.classes:
-                inputs = torch.rand(batch_size, *shape)
-                transform, mask = create_coupling_transform(cls, shape)
-                outputs, logabsdet = transform(inputs)
-                with self.subTest(cls=cls, shape=shape):
-                    self.assert_tensor_is_good(outputs, [batch_size] + shape)
-                    self.assert_tensor_is_good(logabsdet, [batch_size])
-                    self.assertEqual(
-                        outputs[:, mask <= 0, ...], inputs[:, mask <= 0, ...]
-                    )
-
-    def test_inverse_unconstrained(self):
-        for shape in self.shapes:
-            for cls in self.classes:
-                inputs = 3.0 * torch.randn(batch_size, *shape)
-                transform, mask = create_coupling_transform(cls, shape, tails="linear")
-                outputs, logabsdet = transform(inputs)
-                with self.subTest(cls=cls, shape=shape):
-                    self.assert_tensor_is_good(outputs, [batch_size] + shape)
-                    self.assert_tensor_is_good(logabsdet, [batch_size])
-                    self.assertEqual(
-                        outputs[:, mask <= 0, ...], inputs[:, mask <= 0, ...]
-                    )
-
-    def test_forward_inverse_are_consistent(self):
-        for shape in self.shapes:
-            for cls in self.classes:
-                inputs = torch.rand(batch_size, *shape)
-                transform, mask = create_coupling_transform(cls, shape)
-                with self.subTest(cls=cls, shape=shape):
-                    self.eps = 1e-4
-                    self.assert_forward_inverse_are_consistent(transform, inputs)
+            for tails in [None, "linear"]:
+                with self.subTest(shape=shape, tails=tails):
+                    inputs = torch.rand(batch_size, *shape)
+                    flow, _ = create_coupling_transform(shape, tails=tails)
+                    self.checkForwardInverse(flow, inputs)
 
-    def test_forward_inverse_are_consistent_unconstrained(self):
-        self.eps = 1e-5
-        for shape in self.shapes:
-            for cls in self.classes:
-                inputs = 3.0 * torch.randn(batch_size, *shape)
-                transform, mask = create_coupling_transform(cls, shape, tails="linear")
-                with self.subTest(cls=cls, shape=shape):
-                    self.eps = 1e-4
-                    self.assert_forward_inverse_are_consistent(transform, inputs)
-
-    def test_forward_unconditional(self):
+    def test_rqcs_unconditional(self):
         for shape in self.shapes:
-            for cls in self.classes:
+            with self.subTest(shape=shape):
                 inputs = torch.rand(batch_size, *shape)
                 img_shape = shape[1:] if len(shape) > 1 else None
-                transform, mask = create_coupling_transform(
-                    cls, shape, apply_unconditional_transform=True, img_shape=img_shape
+                flow, _ = create_coupling_transform(
+                    shape, apply_unconditional_transform=True, img_shape=img_shape
                 )
-                outputs, logabsdet = transform(inputs)
-                with self.subTest(cls=cls, shape=shape):
-                    self.assert_tensor_is_good(outputs, [batch_size] + shape)
-                    self.assert_tensor_is_good(logabsdet, [batch_size])
-                    self.assertNotEqual(
-                        outputs[:, mask <= 0, ...], inputs[:, mask <= 0, ...]
-                    )
+                self.checkForwardInverse(flow, inputs)
 
 
 if __name__ == "__main__":