Custom Mean And Covariance for Variational Candidate in BNN example

Hi There,

What would be the best way to access the mean and the covariance of the variational posterior distribution in the bayesian_neural_network.py example?

Everything is happening under the hood in model.train_on_batch where the optimizer (no matter which) computes those hyperparameters (parameters of the Gaussian variational candidate) using a descent step after computing a Monte Carlo (Flipout) approximation of the loss function (embedded in the layers of the model + cross-entropy).

This latter MC approximation requires sampling from the variational candidate with the current values of the hyperparameters: That's where I would like to have access to, in order to change those hyperparameters value.

I could do it in a customer optimizer but was wondering if an easier way could be to do the changes right before the sampling step.

Thanks a lot.
Belhal

Start with the bayesian_neural_network.py example and the model.train_on_batch flow described in the issue. Trace where the variational candidate is sampled and how its mean and covariance are represented, then determine the access point needed to change them before sampling. Done means the example can access and modify those parameters without requiring a custom optimizer.