Pyro vs Zuko

[EthanMarx]

Hey @francois-rozet! I have a project that currently builds NF's with pyro. Its my understanding that the NF portion of pyro is no longer maintained.

I was curious to hear more about your plans for zuko. Are you planning on maintaining it in the near term? Is it feature complete with respect to pyro? Im planning on migrating to flows built with zuko, especially since its actively maintained.
Thanks a ton!

[francois-rozet]

Hi @EthanMarx 👋 Yes, I am planning on maintaining Zuko (and its sister Azula) for a while. Zuko has been stable for almost two years now, and there has been virtually no bug reports since then, so it isn't a lot of work 😁

I think Zuko covers all NF-related features in Pyro and nflows, and more. Zuko's interface is inspired by Pyro's, so you shouldn't feel lost. However, Zuko is much more efficient, especially for auto-regressive and coupling flows, which I spent a lot of time optimizing.

I hope you have fun building with Zuko!

[EthanMarx]

Thanks! It took all of ~5 minutes to convert to zuko. The interface is a lot more intuitive and cleaner.

Can you comment on the efficiency improvement? Are you referring to memory, speed or both? Anything to learn from these optimizations?

[francois-rozet]

Most of the efficiency comes from the transformations. Unlike other packages, all transformations in Zuko implement a call_and_ladj method, which applies the transformation and calculates its (log abs det) Jacobian at the same time, which is typically much more efficient. Another major improvement is the inverse of the auto-regressive transform, which does not require any indexing. You can check this discussion to understand how it works.