Refactor flows module

[francois-rozet]

The zuko/flows.py file has become too long. It makes it hard to navigate the code and add new features. I think splitting the module into sub-modules would help greatly. This PR is a first attempt to refactor the flows module.

[francois-rozet]

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.

[simonschnake]

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.

[francois-rozet]

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

[francois-rozet]

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.