Conditional normalizing flows

I have been experimenting with conditional normalizing flows (NFs), using the ‘conditional_spline’ and ‘conditional_spline_autoregressive’ helper functions as examples of element-wise and multivariate conditional NFs. I am still new to Pyro, and I hope you can help answer a few questions:

  1. Previously, I have used spline-based coupling layers for unconditional NFs. However, I couldn’t find their conditional counterparts. Are there plans to add them in the future?

  2. Among the existing conditional flows, which is considered the most expressive for sampling tasks? Is the conditional autoregressive spline that I’m using the best choice?

  3. Is it possible to restrict the support of the transformed conditional distribution? For example, if I work in 2D and both the x-axis and y-axis have hard lower and upper bounds, can these constraints be enforced during training?

Thanks for your help!

  1. afaik nope; look at other libraries like Zuko 1.4.1?
  2. probably hard to say in general?
  3. use transforms like SigmoidTransform and AffineTransform etc?

Thanks @martinjankowiak! On point 3, are there minimal coding examples I can experiment with? Pyro’s NF tutorial didn’t have cases where several transforms are composed together. If there isn’t anything handy, I’ll give this a try soon and write back.

i’d look through the repo for examples of using TransformeDistribution e.g. here