title: How does one use keras add_weight() vars with tensorflow probability distributions?
body: <p>I am creating a new keras layer which accepts a vector of input data and is parameterized by 2 scalars, a mean and standard deviation.  I model the input data as a normal distribution and estimate its mean and variance through gradient descent.  However, when I initialize tfp.Normal(mu, sigma) which mu and sigma are from add_weights() during, build(), the gradients do not propagate through mu and sigma.</p>
  <p>The tensorflow probability documentation states that you can pass in training variables for distribution parameters and backprop through them.  How do I get this to work inside of keras?</p>
  <p>Below is a minimal working example.</p>
  <pre><code>import tensorflow as tf
  import tensorflow_probability as tfp
  import numpy as np
  tfk = tf.keras
  tfkl = tf.keras.layers
  tfd = tfp.distributions
  tfpl = tfp.layers
  EPS = 1e-5
  
  batch_size = 4
  N = 100
  x = np.random.randn(batch_size, N)
  
  class NormalLikelihood(tf.keras.layers.Layer):
      def __init__(self):
          super(NormalLikelihood, self).__init__()
  
      def build(self, input_shape):
          self.mu = self.add_weight(&quot;mean&quot;, shape=[1], initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=1), dtype=tf.float32)
          self.sigma = self.add_weight(&quot;std&quot;, shape=[1], initializer=tf.keras.initializers.RandomUniform(minval=EPS, maxval=5.0, seed=None), constraint=tf.keras.constraints.non_neg(), dtype=tf.float32)
          self.distribution = tfp.distributions.Normal(self.mu[0], self.sigma[0])
  
      def call(self, input):
          r = self.distribution.prob(input)
          r = tf.clip_by_value(r, 1e-3, 1-1e-3)
          return r
  
  input_layer = tf.keras.layers.Input(shape=(100,))
  r = NormalLikelihood()(input_layer)
  r = -tf.reduce_sum(tf.math.log(r))
  model = tf.keras.models.Model(input_layer, r)
  model.add_loss(r)
  model.compile(optimizer='rmsprop', loss=None)
  model.fit(x, y=None)
  </code></pre>
  <p>This code results in builtins.ValueError: No gradients provided for any variable: ['normal_likelihood/mean:0', 'normal_likelihood/std:0'] which is not expected.  Desired behavior would be that ['normal_likelihood/mean:0', 'normal_likelihood/std:0'] have gradients provided for them.</p>
  <p>See the code in google colab: <a href="https://colab.research.google.com/drive/1_u4XTCIH-2qwNSgv9zkZiCG_zeCIEZGp?usp=sharing" rel="nofollow noreferrer">https://colab.research.google.com/drive/1_u4XTCIH-2qwNSgv9zkZiCG_zeCIEZGp?usp=sharing</a></p>