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 showing the benefits of such an architecture in the context of causal generative models, and our source code 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, 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.