Initialization-Dependent Priors in Bayesian NN Layers - Feature Request

There is a decent amount of research using Bayesian NN layers with the prior set to the posterior initialization (or a function thereof): see e.g. Dziugaite & Roy, 2017. I am currently using this method in my research and having to use an ugly workaround and I'm sure others will be doing the same.

Thus I propose a simple change: allow the kernel_prior_fn and bias_prior_fn layer arguments to depend on the posterior initialization chosen (so that a simple lambda function could e.g. take the posterior initialization and set it as the prior). Alternatively (to ensure more straightforward backward compatibility), extra lambda_kernel_prior_fn arguments could be added.

I believe this would only require simple tweaks to the _DenseVariational and _ConvVariational build functions.

Thanks!

Start by reading the _DenseVariational and _ConvVariational build functions, focusing on how posterior initialization and the kernel_prior_fn and bias_prior_fn arguments are currently handled. Determine how initialization-dependent priors should preserve backward compatibility, then verify that both variational layer types can derive priors from their chosen posterior initialization.