Add tutorials using normalizing flows

[francois-rozet]

This PR adds tutorials (SVI and VAE) using normalizing flows. The normalizing flow implementations are loaded from the Zuko library. By the way, I think it could be interesting to add a "compatibility with external tools" section to the examples.

In addition (as discussed in #3297), this PR also updates the normalizing flow tutorial to mention that Pyro's built-in flow API is not developed anymore. I also fixed a mistake ($\det | \frac{\partial f}{\partial x} |$ should be $| \det \frac{\partial f}{\partial x} |$) in the math and cleaned up a bit.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[martinjankowiak]

the two usage tutorials look great but i think it'd be much better if you first made a separate PR for Zuko2Pyro which could sit in distributions or in contrib.zuko. this way this class can be tested and imported into your example notebooks

[francois-rozet]

the two usage tutorials look great but i think it'd be much better if you first made a separate PR for Zuko2Pyro which could sit in distributions or in contrib.zuko. this way this class can be tested and imported into your example notebooks

I didn't know I could do that! That would be much better. In fact, this is actually a torch-to-Pyro class (Zuko uses plain PyTorch distributions), so it could be interesting as part of pyro.distributions.

[martinjankowiak]

@francois-rozet isn't rsample_and_log_prob zuko-specific?

[francois-rozet]

Indeed, but the class still works if dist does not implement rsample_and_log_prob. So it is not actually necessary to import zuko. Do you have a preference for pyro.distributions or contrib.zuko ?

BTW, Transform.call_and_ladj (also Zuko-specific) and Distribution.rsample_and_log_prob could be worth moving upstream (to torch itself), as they allow sampling and scoring without cache (and barely increase API complexity).

[martinjankowiak]

@fritzo do you have a preference between pyro.distributions and contrib.zuko? i do not. the former is probably more discoverable...

[martinjankowiak]

note you can put a test for your contrib code in tests/contrib

[martinjankowiak]

@fritzo do we need to add zuko to EXTRAS_REQUIRE?

[francois-rozet]

@martinjankowiak I have addressed your comments (at least the ones I could). I don't think it is necessary to add zuko in EXTRAS_REQUIRE as even contrib/zuko.py does not rely on it. However, it seems like doctest tries to run the example I added in the docstring and fails because it does not import zuko. I can modify the example to be a code-block, which should not be executed.

[martinjankowiak]

lgtm. @fritzo could you please take a final pass, especially w.r.t. caching logic?

[fritzo]

@francois-rozet thank you for cleaning up our existing tutorials and for adding this new tutorial! I think these will help lots of users provide exposure for Zuko. And thanks for your patience in this very slow holiday review cycle. I think we can push a Pyro 1.9 release soon after this merges.

LGTM, subject to @martinjankowiak's review and a couple nits.

[francois-rozet]

I addressed your comments, and I think the tests have passed (went to bed before the end).

Nope i fucked up my dummy test. Fixing that tomorrow.

[martinjankowiak]

lgtm apart from possibly adding a rsample_and_log_prob test

[francois-rozet]

apart from possibly adding a rsample_and_log_prob test

I added a test already.

[martinjankowiak]

sorry i missed that thanks

[francois-rozet]

🎉

[martinjankowiak]

thanks @francois-rozet !!