Continuous Normalizing Flows (CNFs) and Energy Models

[Zymrael]

A planned feature for the next version of torchdyn is a model zoo containing models such as Stable Neural ODEs, Hamiltonian Neural ODEs, Lagrangian Neural ODEs, as well as CNFs and their many variants.

These will be wrappers to nn.Module interfacing the generic neural network defining the dynamics with the DEFunc class, which handles augmentation, depth-variance, data-control and other auxiliary features that allow for mixing and matching of Neural ODE variants.

The design choice follows from the observation that the models above all require additional steps at each function evaluation of the solver e.g gradient computations for Hamiltonian Neural ODEs, Jacobian / hessian calculations for Lagrangian Neural ODEs, handling log-probability propagation for CNFs. It seems thus natural to include these as wrappers, which can then even be stacked and passed to a DEFunc and then the NeuralDE itself.

We are looking for feedback on the above, as we expect CNFs and the other models to be a popular feature in torchdyn. Check new_release_dev for more information on the WIP version.

Also accepting PRs with proof of concept implementations or tutorials on the topic.

[Zymrael]

v0.2.0 adds CNFs, but I'll leave the discussion open for feedback on the current implementation.

[Zymrael]

Discussion has been moved to Slack, closing the Issue.