Is it possible to fix the mean/var arg in tfpl.IndependentNormal?

Question

I wrote a standard probabilistic neural network in Python with the last layer being tfp.layers.IndependentNormal giving me a normal distribution. However, I just want to train the mean of said distribution, leaving the variance fixed.

Has anyone tried anything similar or has an idea how to do that?

Answer 1

You can do it like this:

model = Sequential([
    ...
    tfpl.DistributionLambda(lambda t: tfd.Independent(tfd.Normal(loc = t, scale = 0.5)))
])

Here you need to set scale into a constant value to keep it fixed.