Tensorflow Probability change of prior in tfp.layers (Possible Issue)

Hello,

My research needs to change the prior distribution parameters to train a Bayesian Neural Network model. More specifically, I want to change the prior Gaussian distribution mu and sigma.

I have read the core codes of Tensorflow Probability and implemented my own code for changing the prior. However, after during a couple of extreme experiments it turned out that the results directly contradicts the theory. Here, the problem is either with the way that I am implementing the prior, or with the Bayesian weight updates which Tensorflow implements or the derivation behind these weight updates.

Now, the easiest way is that I make sure I am implementing the change of prior in my codes correct. Can you please let me know what is the correct way to implement another prior for tfp.layers ?

My code for implementing this is as follow:

`

tfp.layers.Convolution2DReparameterization(
          filters=6, kernel_size=(5,5), padding='same', activation='relu', dtype=‘float64’,
          kernel_divergence_fn=lambda q, p, _: tfp.distributions.kl_divergence(q, p) / num_example,
          kernel_prior_fn= _Prior_fn(mean=initialization[0], scale=initialization[1]),
          bias_posterior_fn=tfp.layers.default_mean_field_normal_fn(),
          bias_prior_fn= _Prior_fn(mean=initialization[2], scale=initialization[3]),
          bias_divergence_fn=lambda q, p, _: tfp.distributions.kl_divergence(q, p) / num_example)

def _Prior_fn(mean, scale,*args, **kwargs):
    d = tfd.Normal(loc=mean, scale=scale)
    def fn(*args, **kwargs):
        return tfd.Independent(d, reinterpreted_batch_ndims=tf.size(d.batch_shape_tensor()))
    return fn

`

**Initialization[0], initialization[1] have the exact shape of the kernel_posterior, and initialization[3], initialization[4] have the exact shape of bias_posterior.

Can you please verify this code, or give me another sample code to implement my own prior to pass to tfp.layers?

Thanks a lot for you help,

No source files or tests are named. Start by reproducing the reported behavior with the provided tfp.layers.Convolution2DReparameterization and _Prior_fn example, then compare the custom prior and divergence configuration with TensorFlow Probability's layer documentation; done means determining whether the implementation or the reported Bayesian update behavior is incorrect.