Add CVAE in Flax

[dirmeier]

Hello,

This PR is a port of @fritzo 's CVAE example. It's not particularly fancy, but I believe helpful nonetheless, since (a) it shows how to use optax.multi_transform to freeze a set of parameters while training (i.e., not apply gradient transforms) and use that with SVI, and (b) it gives another Flax example which I think would be good to have more of.

Hope it's useful.

Cheers,
Simon

[dirmeier]

I think the modelling tests fail, because flax seems to require installation of pyyaml now?

[fehiepsi]

That's right. You can merge master to make tests passed.

Is this ready to review? The overall code looks great to me. Maybe it is helpful to add some README to help other people navigate the folder, point to the original Pyro tutorial to illustrate the data and the model. If you want me to look at some specific points, please let me know.

[dirmeier]

Is this ready to review? The overall code looks great to me. Maybe it is helpful to add some README to help other people navigate the folder, point to the original Pyro tutorial to illustrate the data and the model. If you want me to look at some specific points, please let me know.

Yes, it is ready to review.

I'd like to wrap model and guide as a Flax module like in the original example here, e.g., something like

from flax import linen as nn

class CVAE(nn.Module):
  
  def setup(self):
    self.baseline = BaselineNet()
    self.prior_net = Encoder()
    ...
   
  def model(self, x, y):
    ...

   def guide(self, x, y):
     ...

Do you know if this works somehow? So far I am getting errors with various attempts.

[fehiepsi]

It is tricky to use such pattern. I feel that the current approach is cleaner. No need to mix neural network definition with random variable in the same class.

[fehiepsi]

I think we can remove jit here

[dirmeier]

@fehiepsi could you run the workflows again please? I don't think it fails bcs of my changes..

[fehiepsi]

Thanks for your contribution, @dirmeier! Could you do some extra steps to make this great example more visible for other folks?

• Add a cvae.py file with only docstrings like in your README.md file (you can use this example as a reference). Add a link to github folder of the example (I guess it is https://github.com/pyro-ppl/numpyro/tree/master/examples/cvae-flax)
• Move cvae_predictions.png to this folder, rename it to cvae.png
• Add examples/cvae to the application section of [sphinx]((https://github.com/pyro-ppl/numpyro/blob/master/docs/source/index.rst)

[dirmeier]

All done, let me know if this works for you or you 'd like to see some other changes..
Cheers,
S

[fehiepsi]

Thanks @dirmeier!