code: def test_lower_bound_jensen_shannon(self):
      # Check against numpy implementation.
      log_f = self.optimal_critic
      js_bound, log_f = self.evaluate(
          [mi.lower_bound_jensen_shannon(log_f), log_f]
      )
      # The following numpy softplus is numerically stable when x is large
      # log(1+exp(x)) = log(1+exp(x)) - log(exp(x)) + x = log(1+exp(-x)) + x
      numpy_softplus = lambda x: np.log(
          1 + np.exp(-np.abs(x))
      ) + np.maximum(x, 0)
  
      log_f_diag = np.diag(log_f)
      n = np.float32(log_f.shape[0])
      first_term = np.mean(-numpy_softplus(-log_f_diag))
      second_term = (
          np.sum(numpy_softplus(log_f))
          - np.sum(numpy_softplus(log_f_diag))
      ) / (n * (n - 1.0))
      numpy_implementation = first_term - second_term
      self.assertAllClose(js_bound, numpy_implementation, rtol=1e-5)
  
      # Check the masked against none masked version
      js_bound_1 = mi.lower_bound_jensen_shannon(self.scores)
      positive_mask = np.eye(self.scores.shape[0], self.scores.shape[1])
      js_bound_2 = self.evaluate(
          mi.lower_bound_jensen_shannon(
              self.scores, positive_mask, validate_args=True
          )
      )
      self.assertAllClose(js_bound_1, js_bound_2)
  
      # Check batched against none batched version
      js_bound_3 = self.evaluate(
          mi.lower_bound_jensen_shannon(
              tf.tile(self.scores[tf.newaxis, :, :], [3, 1, 1])
          )
      )
      self.assertAllClose(
          js_bound_3,
          self.evaluate(tf.tile(js_bound_1[tf.newaxis,], [3])),
      )
  
file_path: probability/tensorflow_probability/python/vi/mutual_information_test.py
imports: import numpy
  import scipy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.vi import mutual_information
label_desc: Use image node_id 7 for calling the MutualInformationTest obj's underlying member method code with example usage: obj.test_lower_bound_jensen_shannon() without return types