Gaussian Mixture via TFP layers

Hello TF Community! :slight_smile: First time poster and I’d really appreciate your help!

I’ve been following this excellent blog post about VAE and the use of Tensorflow probability.

However, even after reading the documentation, it is not entirely clear to me what the final probability layer outputs?

Also, I am quite interested if it was possible to use TFP for Gaussian Mixture models that would be integrated with VAE. That is, using prior MixtureSameFamily and for the final layer? It seems there is not much research around this.

Maybe someone could help me out :slightly_smiling_face: