title: Include losses defined in layers using trainable parameters into the final loss for optimization
body: <p>Hi I am trying to create a VariationalDenseLayer where the KL loss is computed in the call function.</p>
  <pre><code>def call(self, inputs, **kwargs):
      kernel_sigma = tf.math.softplus(self.kernel_rho)
      kernel = self.kernel_mu + kernel_sigma * tf.random.normal(self.kernel_mu.shape)
  
      bias_sigma = tf.math.softplus(self.bias_rho)
      bias = self.bias_mu + bias_sigma * tf.random.normal(self.bias_mu.shape)
  
      self.add_loss(self.kl_loss(kernel, self.kernel_mu, kernel_sigma) +
                    self.kl_loss(bias, self.bias_mu, bias_sigma))
  
      outputs = gen_math_ops.MatMul(a = inputs, b = kernel)
  
      if self.use_bias:
          outputs = nn_ops.bias_add(outputs, bias)
  
      if self.activation is not None:
          outputs = self.activation(outputs)
  
      return outputs
  
  
  def kl_loss(self, w, mu, sigma):
      variational_dist = tfp.distributions.Normal(mu, sigma)
      return self.kl_weight * K.sum(variational_dist.log_prob(w) - self.log_prior_prob(w))
  
  def log_prior_prob(self, w):
      comp_1_dist = tfp.distributions.Normal(0.0, self.prior_sigma_1)
      comp_2_dist = tfp.distributions.Normal(0.0, self.prior_sigma_2)
      return K.log(self.prior_pi_1 * comp_1_dist.prob(w) +
                   self.prior_pi_2 * comp_2_dist.prob(w))
      
  </code></pre>
  <p>however, I got following errors when I tried to defined my final loss to be neg_log_likelihhood + kl_loss:</p>
  <blockquote>
  <p>TypeError: An op outside of the function building code is being passed
  a &quot;Graph&quot; tensor. It is possible to have Graph tensors leak out of the
  function building context by including a tf.init_scope in your
  function building code. For example, the following function will fail:
  @tf.function   def has_init_scope():
  my_constant = tf.constant(1.)
  with tf.init_scope():
  added = my_constant * 2 The graph tensor has name: dense_flipout/divergence_kernel:0</p>
  </blockquote>
  <p>If I added</p>
  <pre><code>tf.compat.v1.disable_eager_execution()
  </code></pre>
  <p>I'd got follow errors.  However, it is still not clear to me what I did wrong.</p>
  <blockquote>
  <p>loss = neg_log_likelihood + kl * kl_weight   File &quot;/d/dev01/chungh/lib2/anaconda3-2020.11/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py&quot;, line 1266, in r_binary_op_wrapper
  y, x = maybe_promote_tensors(y, x)   File &quot;/d/dev01/chungh/lib2/anaconda3-2020.11/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py&quot;, line 1202, in maybe_promote_tensors
  ops.convert_to_tensor(tensor, dtype, name=&quot;x&quot;))   File &quot;/d/dev01/chungh/lib2/anaconda3-2020.11/lib/python3.8/site-packages/tensorflow/python/profiler/trace.py&quot;, line 163, in wrapped
  return func(*args, **kwargs)   File &quot;/d/dev01/chungh/lib2/anaconda3-2020.11/lib/python3.8/site-packages/tensorflow/python/framework/ops.py&quot;, line 1566, in convert_to_tensor
  ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)   File
  &quot;/d/dev01/chungh/lib2/anaconda3-2020.11/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py&quot;,
  line 339, in _constant_tensor_conversion_function
  return constant(v, dtype=dtype, name=name)   File &quot;/d/dev01/chungh/lib2/anaconda3-2020.11/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py&quot;,
  line 264, in constant
  return _constant_impl(value, dtype, shape, name, verify_shape=False,   File
  &quot;/d/dev01/chungh/lib2/anaconda3-2020.11/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py&quot;,
  line 281, in _constant_impl
  tensor_util.make_tensor_proto(   File &quot;/d/dev01/chungh/lib2/anaconda3-2020.11/lib/python3.8/site-packages/tensorflow/python/framework/tensor_util.py&quot;,
  line 457, in make_tensor_proto
  _AssertCompatible(values, dtype)   File &quot;/d/dev01/chungh/lib2/anaconda3-2020.11/lib/python3.8/site-packages/tensorflow/python/framework/tensor_util.py&quot;,
  line 336, in _AssertCompatible
  raise TypeError(&quot;Expected %s, got %s of type '%s' instead.&quot; % TypeError: Expected float32, got &lt;function neg_log_likelihood at
  0x2aab06930b80&gt; of type 'function' instead.</p>
  </blockquote>
  <p>The code to compute the final loss is:</p>
  <pre><code>def neg_log_likelihood(y_obs, y_pred, sigma=noise):
      dist = tfp.distributions.Normal(loc=y_pred, scale=sigma)
      return K.sum(-dist.log_prob(y_obs))
  
  kl = sum(model.losses)
  loss = kl * kl_weight
  
  model.compile(loss=loss, optimizer=optimizers.Adam(lr=0.08), metrics=['mse'])
  </code></pre>
  <p>Any helps and pointers are highly appreciated!</p>