Add checkerboard inversion

normalizing_flows/bijections/finite/autoregressive/conditioning/coupling_masks.py

@@ -78,12 +78,13 @@ class Checkerboard(Coupling):
     Checkerboard coupling for image data.
     """
 
-    def __init__(self, event_shape, resolution: int = 2):
+    def __init__(self, event_shape, resolution: int = 2, invert: bool = False):
         """
         :param event_shape: image shape with the form (n_channels, width, height). Note: width and height must be equal
         and a power of two.
         :param resolution: resolution of the checkerboard along one axis - the number of squares. Must be a power of two
         and smaller than image width.
+        :param invert: invert the checkerboard mask.
         """
         n_channels, width, _ = event_shape
         assert width % resolution == 0
@@ -93,6 +94,8 @@ def __init__(self, event_shape, resolution: int = 2):
         a = torch.tensor([[1, 0] * half_resolution, [0, 1] * half_resolution] * half_resolution)
         mask = torch.kron(a, torch.ones((square_side_length, square_side_length)))
         mask = mask.bool()
+        if invert:
+            mask = ~mask
         super().__init__(event_shape, mask)
 
 

test/test_checkerboard_coupling.py

@@ -44,3 +44,19 @@ def test_checkerboard_medium():
         ], dtype=torch.bool)[None].repeat(3, 1, 1)
     )
     assert torch.allclose(coupling.target_mask, ~coupling.source_mask)
+
+
+def test_checkerboard_small_inverted():
+    torch.manual_seed(0)
+    image_shape = (3, 4, 4)
+    coupling = Checkerboard(image_shape, resolution=2, invert=True)
+    assert torch.allclose(
+        coupling.source_mask,
+        ~torch.tensor([
+            [1, 1, 0, 0],
+            [1, 1, 0, 0],
+            [0, 0, 1, 1],
+            [0, 0, 1, 1],
+        ], dtype=torch.bool)[None].repeat(3, 1, 1)
+    )
+    assert torch.allclose(coupling.target_mask, ~coupling.source_mask)