Hey @simonschnake, I have updated the README in this PR, do you think the new introduction is good? Feel free to make comments on the PR, if you find the time.

Hey @francois-rozet,
I think the README very well written and good understandable.
Are you planing to switch the naming from TransformModule to TranformFactory?

I would change one sentence in the README:
“This design allows for flows that behave like distributions while retaining the benefits of Module.“
->
“This design enables flows to act like distributions while maintaining the feature of trainable weights inherent in a Module.”
That is a bit more precise.

Another thing is spot was

flow = FlowModule(
    transforms=[
        MaskedAutoregressiveTransform(3, 5, hidden_features=[128] * 3),
        Unconditional(PermutationTransform, torch.randperm(3), buffer=True),
        MaskedAutoregressiveTransform(3, 5, hidden_features=[128] * 3),
    ],
    base=Unconditional(
        DiagNormal,
        torch.zeros(3),
        torch.ones(3),
        buffer=True,
    ),
)

Here you take as an example, to make the second layer a PermutationTransform. One problem could be, that a new reader would think that this is the way to include permutations in a MaskedAutoregressiveFlow and not that this is a build in functionality of the MaskedAutoregressiveTransform. I would propose to pick a different example or include a sentence which points to that.

These are just small things. Overall, I think it is a very solid piece of work.

Thanks! I addressed your comments for the README. I will merge soon.

The new API documentation is up at https://zuko.readthedocs.io/en/stable/api/zuko.flows.html 🔥 I think it is much more readable than before. I also added a few warnings about invertibility. Next step are tutorials.