Hi,

I’m trying to fit a radial BNN posterior variational approximation as per this paper.

However, since I’ll be training a BNN, I don’t want to have to write a custom guide and define this variational approximation for all of my layers, and so was trying to implement a custom AutoGuide which automatically puts a radial BNN approximation on all of my weights.

The radial approximation is defined as follows:

where I just need to sample all epsilon_MFVI from an independent standard normal distribution, normalize them, and multiply them by r, which is a scalar sampled from a standard normal.

How could I go about implementing this custom AutoGuide?

Is there a smarter way of implementing this variational approximation?

Thanks in advance.