But a VAE is based on probability. I dont see this in the code you provided. A softmax output cannot be interpreted as the uncertainty of a model. You can only achieve this with having a distribution over the weights of a layer. And this is where tf probability comes into place or the monto carlo dropout during inference mechanism that is proven to work.Anyway the methodology you provided looks very interesting. I do not want to offend you by any means.