How to instantiate variable-containing distribution inside tfp.layers.DistributionLambda

Question

I\'m trying to build a tf.keras.Sequential model using tfp.layers.DistributionLambda. I\'m following the DistributionLambda example, but would want