Refactor two moons normalizing flow figure and add Continuous Normalizing Flow example

[vz415]

The normalizing flow chapter (chpt. 22 in book 2) does not have an example for continuous normalizing flows (CNFs), which uses differential equations to map the flow. I believe this is an important component to include in the book given its intersection between probabilistic and physical modeling using differenital equations.

I have created a two moons Jax demo based off of Patrick Kidger's demo using his diffrax differential equation library. My demo (found here) provides and example and includes some math theory from the book to help readers understand the code. My demo does not use distrax but instead opts for equinox to use MLP-like layers (non-invertible). It was easier to reverse layers with equinox than with haiku and CNF does not need invertible layers due to differential equation theory (explained in the notebook).

Additionally, the current two moons figure could be improved by showing the process of data going from a Gaussian base to data distribution, p(x). I think Figure 5b from the Neural Ordinary Differential Equations paper provides a good example of a flow going from base to data distribution. I'd like to improve this figure by adding sampling from different layers of the current NSF example, and compare to the CNF.

Figure 5b from Neural ODE paper.

Proposal: Replace Fig 22.8 in normalizing flow notebook with two moons generated from both NSF and CNF to replicate Fig 5b from Neural ODE paper.

I can either make this one complete notebook for the figure or leave two separate notebooks. I'm also happy to update the notebook to latexify it if needed. Happy to hear feedback.

[murphyk]

Adding CNF would be great. Also I agree that the current fig 22.8 showing neural spline flow is not a very informative figure. Please feel free to open a PR with a new notebook. I think a single one is fine for now (we may split later if it's too big.)