Repository: probabilists/zuko

Issue:
Add support for custom adjacency matrices in autoregressive models
### Description

I would like to enable support for autoregressive models (that is, MAF, NSF, and NAF) to use a specific adjacency matrix (rather than an ordering), which could be sparser than the triangular matrix obtained from the ordering. In case an adjacency matrix is provided, all layers would use that adjacency matrix for the masking, instead of using random permutations between layers.

The motivation for this feature is that sometimes adding extra structure to the normalizing flows can be really useful in practice. Past year, we [published a paper](https://arxiv.org/abs/2306.05415) showing the benefits of such an architecture in the context of causal generative models, and our [source code](https://github.com/psanch21/causal-flows) actually adapted Zuko to use adjacency matrix for this purpose.

Since the changes required are not many, I thought I could make a pull request with this feature so that we can reuse Zuko for future projects with much more ease (since the version from our repository is based un Zuko v0.2.2).

### Implementation

The implementation would follow that from [our repository](https://github.com/psanch21/causal-flows), and it basically consists in adding the adjacency matrix as another parameter of the constructor, and use it if provided, rather than building the matrix from the ordering.

I have a pull-request ready, which I will link to this issue after submitting it.

### Alternatives

I did not think of any alternatives, since: i) the functionality is already implemented, the only real barrier is that it was not a parameter of the constructor; and ii) we have a working and tested implementation.

Discussion:
@francois-rozet: Hello @adrianjav, thank you for the feature request and PR! I think this is a nice and sensible feature to add to Zuko. I will review the PR soon.

---

On another note, I passed by your poster at NeurIPS last year and enjoyed your paper! I had no idea your implementation was based on Zuko :sweat_smile: I am grateful that you found Zuko useful for your research, but I have to say I am saddened that you did not mention it in the manuscript or the README.

In fact, I notice that you copy-pasted a large portion of Zuko's code in your repository without any credits, which is explicitly prohibited by our license.

> The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

I really don't want to pass as a prick, but could you please address this?@adrianjav: Hi @francois-rozet, thanks for taking a look at the PR, and I am glad that you agree this is sensible feature to have.

---

I am glad to know you liked our work 😃 

You are completely right and entitled to ask us to credit your work. _I am really sorry_ we didn't do it earlier and, to be honest, I have not stopped to realize that we did not credit your work until you have mentioned it. I will fix it and credit Zuko both in the repository (asap) and in the latest arxiv version (soon). I will contact you privately to double-check with you.

PS: I see that you are nowadays working with diffusion models, but I'd be happy to meet and talk science if you ever feel like it!@francois-rozet: Thank you for acting quickly!

> PS: I see that you are nowadays working with diffusion models, but I'd be happy to meet and talk science if you ever feel like it!

I'd be happy to! My academic e-mail is francois.rozet@uliege.be.

Solution patch:
diff --git a/zuko/flows/autoregressive.py b/zuko/flows/autoregressive.py
index 5d878dd..a215c34 100644
--- a/zuko/flows/autoregressive.py
+++ b/zuko/flows/autoregressive.py
@@ -9,7 +9,7 @@ import torch
 
 from functools import partial
 from math import ceil, prod
-from torch import LongTensor, Size, Tensor
+from torch import BoolTensor, LongTensor, Size, Tensor
 from torch.distributions import Transform
 from typing import Callable, Sequence
 
@@ -38,7 +38,10 @@ class MaskedAutoregressiveTransform(LazyTransform):
         passes: The number of sequential passes for the inverse transformation. If
             :py:`None`, use the number of features instead, making the transformation
             fully autoregressive. Coupling corresponds to :py:`passes=2`.
-        order: The feature ordering. If :py:`None`, use :py:`range(features)` instead.
+        order: A feature ordering. If :py:`None`, use :py:`range(features)` instead.
+        adjacency: An adjacency matrix describing the transformation graph. If
+            `adjacency` is provided, `order` is ignored and `passes` is replaced by the
+            diameter of the graph.
         univariate: The univariate transformation constructor.
         shapes: The shapes of the univariate transformation parameters.
         kwargs: Keyword arguments passed to :class:`zuko.nn.MaskedMLP`.
@@ -72,6 +75,7 @@ class MaskedAutoregressiveTransform(LazyTransform):
         context: int = 0,
         passes: int = None,
         order: LongTensor = None,
+        adjacency: BoolTensor = None,
         *args,
         **kwargs,
     ) -> LazyTransform:
@@ -86,6 +90,7 @@ class MaskedAutoregressiveTransform(LazyTransform):
         context: int = 0,
         passes: int = None,
         order: LongTensor = None,
+        adjacency: BoolTensor = None,
         univariate: Callable[..., Transform] = MonotonicAffineTransform,
         shapes: Sequence[Size] = ((), ()),
         **kwargs,
@@ -100,36 +105,97 @@ class MaskedAutoregressiveTransform(LazyTransform):
         # Adjacency
         self.register_buffer('order', None)
 
-        if passes is None:
-            passes = features
+        if adjacency is None:
+            if passes is None:
+                passes = features
 
-        if order is None:
-            order = torch.arange(features)
+            if order is None:
+                order = torch.arange(features)
+            else:
+                order = torch.as_tensor(order, dtype=int)
+
+            assert order.ndim == 1, "'order' should be a vector."
+            assert order.shape[0] == features, f"'order' should have {features} elements."
+
+            self.passes = min(max(passes, 1), features)
+            self.order = torch.div(order, ceil(features / self.passes), rounding_mode='floor')
+
+            adjacency = self.order[:, None] > self.order
         else:
-            order = torch.as_tensor(order)
+            adjacency = torch.as_tensor(adjacency, dtype=bool)
 
-        self.passes = min(max(passes, 1), features)
-        self.order = torch.div(order, ceil(features / self.passes), rounding_mode='floor')
+            assert adjacency.ndim == 2, "'adjacency' should be a matrix."
+            assert adjacency.shape[0] == features, f"'adjacency' should have {features} rows."
+            assert adjacency.shape[1] == features, f"'adjacency' should have {features} columns."
+            assert adjacency.diag().all(), "'adjacency' should have ones on the diagonal."
 
-        in_order = torch.cat((self.order, torch.full((context,), -1)))
-        out_order = torch.repeat_interleave(self.order, self.total)
-        adjacency = out_order[:, None] > in_order
+            adjacency = adjacency * ~torch.eye(features, dtype=bool)
+
+            self.passes = self._dag_diameter(adjacency)
+
+        if context > 0:
+            adjacency = torch.cat(
+                (adjacency, torch.ones((features, context), dtype=bool)),
+                dim=1,
+            )
+
+        adjacency = torch.repeat_interleave(adjacency, repeats=self.total, dim=0)
 
         # Hyper network
         self.hyper = MaskedMLP(adjacency, **kwargs)
 
+    @staticmethod
+    def _dag_diameter(adjacency: BoolTensor) -> int:
+        r"""Returns the diameter of a directed acyclic graph.
+
+        If the graph contains cycles, this function raises an error.
+
+        Credits:
+            This code is adapted from :func:`networkx.topological_generations`.
+
+        Arguments:
+            adjacency: An adjacency matrix representing a directed graph.
+
+        Returns:
+            The diameter of the graph.
+        """
+
+        all_generations = []
+        indegree = adjacency.sum(dim=1).tolist()
+        zero_indegree = [n for n, d in enumerate(indegree) if d == 0]
+        while zero_indegree:
+            this_generation, zero_indegree = zero_indegree, []
+            for node in this_generation:
+                for child in adjacency[:, node].nonzero():
+                    child = child.item()
+                    indegree[child] -= 1
+                    if indegree[child] == 0:
+                        zero_indegree.append(child)
+            all_generations.append(this_generation)
+
+        assert all(d == 0 for d in indegree), "The graph contains cycles."
+
+        return len(all_generations)
+
     def extra_repr(self) -> str:
         base = self.univariate(*map(torch.randn, self.shapes))
-        order = self.order.tolist()
 
-        if len(order) > 10:
-            order = order[:5] + [...] + order[-5:]
-            order = str(order).replace('Ellipsis', '...')
+        if self.order is None:
+            return '\n'.join([
+                f'(base): {base}',
+                f'(passes): {self.passes}',
+            ])
+        else:
+            order = self.order.tolist()
+
+            if len(order) > 10:
+                order = order[:5] + [...] + order[-5:]
+                order = str(order).replace('Ellipsis', '...')
 
-        return '\n'.join([
-            f'(base): {base}',
-            f'(order): {order}',
-        ])
+            return '\n'.join([
+                f'(base): {base}',
+                f'(order): {order}',
+            ])
 
     def meta(self, c: Tensor, x: Tensor) -> Transform:
         if c is not None:
diff --git a/zuko/flows/coupling.py b/zuko/flows/coupling.py
index 9dea154..46a519e 100644
--- a/zuko/flows/coupling.py
+++ b/zuko/flows/coupling.py
@@ -93,16 +93,22 @@ class GeneralCouplingTransform(LazyTransform):
         self.total = sum(prod(s) for s in shapes)
 
         # Mask
-        self.register_buffer('mask', None)
-
         if mask is None:
-            self.mask = torch.arange(features) % 2 == 1
+            mask = torch.arange(features) % 2 == 1
         else:
-            self.mask = mask
+            mask = torch.as_tensor(mask, dtype=bool)
+
+        assert mask.ndim == 1, "'mask' should be a vector."
+        assert mask.shape[0] == features, f"'mask' should have {features} elements."
 
-        features_a = self.mask.sum().item()
+        features_a = mask.sum().item()
         features_b = features - features_a
 
+        assert features_a > 0
+        assert features_b > 0
+
+        self.register_buffer('mask', mask)
+
         # Hyper network
         self.hyper = MLP(features_a + context, features_b * self.total, **kwargs)
 
diff --git a/zuko/transforms.py b/zuko/transforms.py
index d37aec5..d24ad59 100644
--- a/zuko/transforms.py
+++ b/zuko/transforms.py
@@ -981,7 +981,9 @@ class CouplingTransform(Transform):
 
     Arguments:
         meta: A function which returns a transformation :math:`f` given :math:`x_a`.
-        mask: A coupling mask defining the split :math:`x \to (x_a, x_b)`.
+        mask: A coupling mask defining the split :math:`x \to (x_a, x_b)`. Ones
+            correspond to the constant split :math:`x_a` and zeros to the transformed
+            split :math:`x_b`.
     """
 
     domain = constraints.real_vector


Test patch:
diff --git a/tests/test_flows.py b/tests/test_flows.py
index 4d7d44d..fde7087 100644
--- a/tests/test_flows.py
+++ b/tests/test_flows.py
@@ -71,26 +71,37 @@ def test_flows(tmp_path: Path, F: callable):
 
 
 def test_triangular_transforms():
+    order = torch.randperm(5)
+
+    adjacency = torch.rand((5, 5)) < 0.25
+    adjacency = adjacency + torch.eye(5, dtype=bool)
+    adjacency = torch.tril(adjacency)
+    adjacency[1, 0] = True
+    adjacency = adjacency[order, :][:, order]
+
     Ts = [
         ElementWiseTransform,
         GeneralCouplingTransform,
+        partial(GeneralCouplingTransform, mask=order % 2),
         MaskedAutoregressiveTransform,
         partial(MaskedAutoregressiveTransform, passes=2),
+        partial(MaskedAutoregressiveTransform, order=order),
+        partial(MaskedAutoregressiveTransform, adjacency=adjacency),
     ]
 
     for T in Ts:
         # Without context
-        t = T(3)
-        x = randn(3)
+        t = T(5)
+        x = randn(64, 5)
         y = t()(x)
 
-        assert y.shape == x.shape, t
-        assert y.requires_grad, t
+        assert y.shape == x.shape, T
+        assert y.requires_grad, T
         assert torch.allclose(t().inv(y), x, atol=1e-4), T
 
         # With context
-        t = T(3, 5)
-        x, c = randn(256, 3), randn(5)
+        t = T(5, 7)
+        x, c = randn(64, 5), randn(7)
         y = t(c)(x)
 
         assert y.shape == x.shape, T
@@ -98,8 +109,8 @@ def test_triangular_transforms():
         assert torch.allclose(t(c).inv(y), x, atol=1e-4), T
 
         # Jacobian
-        t = T(7)
-        x = randn(7)
+        t = T(5)
+        x = randn(5)
         y = t()(x)
 
         J = torch.autograd.functional.jacobian(t(), x)
@@ -107,3 +118,38 @@ def test_triangular_transforms():
 
         assert torch.allclose(t().log_abs_det_jacobian(x, y), ladj, atol=1e-4), T
         assert torch.allclose(J.diag().abs().log().sum(), ladj, atol=1e-4), T
+
+
+def test_adjacency_matrix():
+    T = MaskedAutoregressiveTransform
+
+    # With a valid adjacency matrix
+    order = torch.randperm(5)
+
+    adjacency = torch.rand((5, 5)) < 0.25
+    adjacency = adjacency + torch.eye(5, dtype=bool)
+    adjacency = torch.tril(adjacency)
+    adjacency[1, 0] = True
+    adjacency = adjacency[order, :][:, order]
+
+    t = T(5, adjacency=adjacency)
+    x = randn(5)
+
+    J = torch.autograd.functional.jacobian(t(), x)
+
+    assert (J[~adjacency] == 0).all()
+
+    # With False in the diagonal
+    adjacency_invalid = adjacency.clone()
+    adjacency_invalid[0, 0] = False
+
+    with pytest.raises(AssertionError, match="'adjacency' should have ones on the diagonal."):
+        t = T(5, adjacency=adjacency_invalid)
+
+    # With cycles
+    adjacency_invalid = adjacency.clone()
+    adjacency_invalid[0, 1] = True
+    adjacency_invalid[1, 0] = True
+
+    with pytest.raises(AssertionError, match="The graph contains cycles."):
+        t = T(5, adjacency=adjacency_invalid)