Trainable priors in Conv2DFlipout layer

Hi all,

I am relatively new to tfp, possibly I have a rather basic doubt. Default Conv2DFlipout layer looks like this

    tfp.layers.Convolution2DFlipout(
        .
        .
        kernel_posterior_fn=tfp_layers_util.default_mean_field_normal_fn(),
        kernel_posterior_tensor_fn=(lambda d: d.sample()),
        kernel_prior_fn=tfp.layers.default_multivariate_normal_fn,
        kernel_divergence_fn=(lambda q, p, ignore: kl_lib.kl_divergence(q, p)), 
        .
        .
        **kwargs
    )

Can someone tell how can I involve kernel and bias priors in training. My intention is to use Empirical Bayes, in very similar lines as given in this blog by the tfp team. I would wish to train my prior mean and keep the prior standard deviation constant.

The default multivariate_normal_fn looks like this,

tfp.layers.default_multivariate_normal_fn(dtype, shape, name, trainable, add_variable_fn)

It has trainable as one of the boolean argument, if I set it True, it throws following error. (I was under the impression that it takes remaining arguments itself as in usual case and overwrites the trainable argument)

    tfp.layers.Convolution2DFlipout(
        kernel_prior_fn = tfp_layers_util.default_multivariate_normal_fn(trainable=True)
    )


    TypeError: default_multivariate_normal_fn() missing 4 required positional arguments: 'dtype', 'shape', 'name', and 'add_variable_fn'

Can someone help me solve the issue?

Start by reading the Conv2DFlipout layer API and the default_multivariate_normal_fn signature, then compare the linked probabilistic-layers blog example. The issue does not identify a file or test; done would require a decided, supported way to make prior means trainable while keeping prior standard deviations fixed, along with validation of that behavior.