This project uses the Burn library to learn the the Scikit-Learn Swiss Roll function. The image below shows the true (orange) and learned (blue) distributions:
The VAE is conditioned on radius from center of spiral. As a result, the trained model generates points at a provided radius (burn_vae.generate(r, n_points) using the python bindings).