I know it is a bit late to respond to this post. But I am going to do it anyway.
I am actually looking at the course you mentioned currently and needless to say, I learned a lot about using TFP. However, for beginners, some of the stuff could have been explained a bit more. For example, in the 1st Week's reading on VAE, in the section "A LOWER-VARIANCE ESTIMATOR OF ELBO", it is mentioned that E_{Z ~ q(z|x)}[log p(x|z)] can be estimated by MCMC. But in the 4th week's tutorial, showing how to maximize ELBO, the tutor grossed over by taking just 1 sample. When I tried to actually take multiple samples and estimate this reconstruction error, I couldn't do it. The program is throwing errors because of the error in tensor shapes. I am desperately trying to fix it. So far couldn't.
Also, the decoder distribution is taken as IndependentBernoulli(), but if I want a gaussian distribution instead, that is,
if we assume that p(x|z) ~ N(mu(z), sigma**2I), ie the sigma is learnable but a shared parameter, independent of z, how to do I go about implementing it?
But overall, I learned many things.
Thanks.