title: How to calculate the gradient of the Kullback-Leibler divergence of two tensorflow-probability distributions with respect to the distribution's mean?
body: <p>In tensorflow-2.0, I am trying to create a <code>keras.layers.Layer</code> which outputs the Kullback-Leibler (KL) divergence between two <code>tensorflow_probability.distributions</code>. I would like to calculate the gradient of the output (i.e. the KL divergence) with respect to the mean value of one of the <code>tensorflow_probability.distributions</code>.</p>
  
  <p>In all my attempts so far, the resulting gradients are <code>0</code>, unfortunately.</p>
  
  <p>I tried implementing a minimal example shown below. I was wondering if the problems might have to do with the eager execution mode of <code>tf 2</code>, as I know of a similar approach that worked in <code>tf 1</code>, where eager execution is disabled by default.</p>
  
  <p>This is the minimal example I tried:</p>
  
  <pre class="lang-py prettyprint-override"><code>import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp
  from tensorflow.keras.models import Model
  from tensorflow.keras.layers import Layer,Input
  
  # 1 Define Layer
  
  class test_layer(Layer):
  
      def __init__(self, **kwargs):
          super(test_layer, self).__init__(**kwargs)
  
      def build(self, input_shape):
          self.mean_W = self.add_weight('mean_W',trainable=True)
  
          self.kernel_dist = tfp.distributions.MultivariateNormalDiag(
              loc=self.mean_W,
              scale_diag=(1.,)
          )
          super(test_layer, self).build(input_shape)
  
      def call(self,x):
          return tfp.distributions.kl_divergence(
              self.kernel_dist,
              tfp.distributions.MultivariateNormalDiag(
                  loc=self.mean_W*0.,
                  scale_diag=(1.,)
              )
          )
  
  # 2 Create model
  
  x = Input(shape=(3,))
  fx = test_layer()(x)
  test_model = Model(name='test_random', inputs=[x], outputs=[fx])
  
  
  # 3 Calculate gradient
  
  print('\n\n\nCalculating gradients: ')
  
  # example data, only used as a dummy
  x_data = np.random.rand(99,3).astype(np.float32)
  
  for x_now in np.split(x_data,3):
  #     print(x_now.shape)
      with tf.GradientTape() as tape:
          fx_now = test_model(x_now)
          grads = tape.gradient(
              fx_now,
              test_model.trainable_variables,
          )
          print('\nKL-Divergence: ', fx_now, '\nGradient: ',grads,'\n')
  
  print(test_model.summary())
  </code></pre>
  
  <p>The output of the code above is</p>
  
  <pre><code>Calculating gradients: 
  
  KL-Divergence:  tf.Tensor(0.0029436834, shape=(), dtype=float32) 
  Gradient:  [&lt;tf.Tensor: id=237, shape=(), dtype=float32, numpy=0.0&gt;] 
  
  
  KL-Divergence:  tf.Tensor(0.0029436834, shape=(), dtype=float32) 
  Gradient:  [&lt;tf.Tensor: id=358, shape=(), dtype=float32, numpy=0.0&gt;] 
  
  
  KL-Divergence:  tf.Tensor(0.0029436834, shape=(), dtype=float32) 
  Gradient:  [&lt;tf.Tensor: id=479, shape=(), dtype=float32, numpy=0.0&gt;] 
  
  Model: "test_random"
  _________________________________________________________________
  Layer (type)                 Output Shape              Param #   
  =================================================================
  input_1 (InputLayer)         [(None, 3)]               0         
  _________________________________________________________________
  test_layer_3 (test_layer)    ()                        1         
  =================================================================
  Total params: 1
  Trainable params: 1
  Non-trainable params: 0
  _________________________________________________________________
  None
  </code></pre>
  
  <p>The KL divergence is calculated correcly, but the resulting gradient is <code>0</code>. What would be a correct way to obtain the gradients?</p>