Repository: probabilists/zuko

Issue:
Add support for custom conditional adjacency
### Description
Zuko currently supports structured dependencies via an adjacency matrix in MaskedAutoregressiveTransform ([#54](https://github.com/probabilists/zuko/pull/54)). This functionality is valuable for leveraging structural inductive biases.

However, the context adjacency—modulating how context conditions each variable—is fixed to an all-ones matrix, implying full context-to-variable conditioning. In [our recent work](https://arxiv.org/abs/2503.15114), we show that structured context adjacencies can be really useful. We propose extending the current implementation to support user-defined context adjacency matrices.

If this aligns with Zuko’s design goals, I would be glad to contribute the implementation via a pull request. We have already implemented this internally for our paper, though the code is not yet public.

### Implementation
Currently, the adjacency matrix is (D × D), with an implicit (C × D) all-ones matrix added when context is used. We propose accepting an extended (D × [D + C]) matrix, where the last C columns specify context-to-variable dependencies. The original behavior is preserved when the matrix is (D × D).

The modification is minimal, localized to MaskedAutoregressiveTransform, and fully backward-compatible.

### Alternatives
No alternatives are currently considered.

Hints from issue comments:
@francois-rozet: Hi @aalmodovares 👋 We will gladly accept a PR for this!

Solution patch:
diff --git a/zuko/flows/autoregressive.py b/zuko/flows/autoregressive.py
index 947b8d0..6b162fd 100644
--- a/zuko/flows/autoregressive.py
+++ b/zuko/flows/autoregressive.py
@@ -40,7 +40,9 @@ class MaskedAutoregressiveTransform(LazyTransform):
         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.
+            diameter of the graph. Its shape must be either `(features, features)`
+            or `(features, features + context)`. If the shape includes context, the rightmost
+            `context` columns define connections to the conditioning variables.
         univariate: The univariate transformation constructor.
         shapes: The shapes of the univariate transformation parameters.
         kwargs: Keyword arguments passed to :class:`zuko.nn.MaskedMLP`.
@@ -120,12 +122,19 @@ class MaskedAutoregressiveTransform(LazyTransform):
             self.order = torch.div(order, ceil(features / self.passes), rounding_mode="floor")
 
             adjacency = self.order[:, None] > self.order
+            adjacency_context = None
         else:
             adjacency = torch.as_tensor(adjacency, dtype=bool)
 
             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.shape[1] == features or adjacency.shape[1] == features + context, (
+                f"'adjacency' should have {features} or {features + context} columns."
+            )
+
+            adjacency_context = adjacency[:, features:] if adjacency.shape[1] > features else None
+            adjacency = adjacency[:, :features]
+
             assert adjacency.diag().all(), "'adjacency' should have ones on the diagonal."
 
             adjacency = adjacency * ~torch.eye(features, dtype=bool)
@@ -133,10 +142,9 @@ class MaskedAutoregressiveTransform(LazyTransform):
             self.passes = self._dag_diameter(adjacency)
 
         if context > 0:
-            adjacency = torch.cat(
-                (adjacency, torch.ones((features, context), dtype=bool)),
-                dim=1,
-            )
+            if adjacency_context is None:
+                adjacency_context = torch.ones((features, context), dtype=bool)
+            adjacency = torch.cat((adjacency, adjacency_context), dim=1)
 
         adjacency = torch.repeat_interleave(adjacency, repeats=self.total, dim=0)
 


Test patch:
diff --git a/tests/test_flows.py b/tests/test_flows.py
index 306193a..04fe959 100644
--- a/tests/test_flows.py
+++ b/tests/test_flows.py
@@ -153,3 +153,42 @@ def test_adjacency_matrix():
 
     with pytest.raises(AssertionError, match="The graph contains cycles."):
         t = T(5, adjacency=adjacency_invalid)
+
+
+def test_context_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 = adjacency[order, :][:, order]
+
+    # context adjacency
+    adjacency_context = torch.rand((5, 2)) < 0.25
+    adjacency_valid = torch.cat((adjacency, adjacency_context), dim=1)
+
+    t = T(features=5, context=2, adjacency=adjacency_valid)
+
+    x, c = randn(5), randn(2)
+    y = t(c)(x)
+
+    assert y.shape == x.shape
+    assert y.requires_grad
+    assert torch.allclose(t(c).inv(y), x, atol=1e-4)
+
+    J = torch.autograd.functional.jacobian(t(c), x)
+
+    assert (J[~adjacency] == 0).all()
+
+    ladj = torch.linalg.slogdet(J).logabsdet
+
+    assert torch.allclose(t(c).log_abs_det_jacobian(x, y), ladj, atol=1e-4)
+    assert torch.allclose(J.diag().abs().log().sum(), ladj, atol=1e-4)
+
+    adjacency_invalid_context = torch.rand((5, 1)) < 0.25
+    adjacency_invalid = torch.cat((adjacency, adjacency_invalid_context), dim=1)
+
+    with pytest.raises(AssertionError, match="'adjacency' should have 5 or 7 columns."):
+        t = T(features=5, context=2, adjacency=adjacency_invalid)