code: def test_masked_logmeanexp(self):
      # test1: compare against numpy/scipy implementation.
      masked_scores = self.scores
      num_masked_ele = np.sum(masked_scores > 0.0)
      masked_scores[masked_scores <= 0.0] = -np.inf
      numpy_impl = np.float32(
          scipy.special.logsumexp(masked_scores)
          - np.log(num_masked_ele)
      )
      result_0d = mi._masked_logmeanexp(
          self.scores, self.scores > 0, axis=None
      )
      self.assertAllClose(self.evaluate(result_0d), numpy_impl)
  
      # test2: test against results from composition of numpy functions.
      scores_2 = np.array(
          [[2.0, 5.0, -1e-3], [-1073.0, 4.2, -4.0]], dtype=np.float32
      )
      result_empty_sum = mi._masked_logmeanexp(
          scores_2, scores_2 < 0.0, axis=None
      )
      numpy_result = np.log(np.mean(np.exp(scores_2[scores_2 < 0.0])))
      self.assertAllClose(
          self.evaluate(result_empty_sum),
          numpy_result.astype(np.float32),
      )
  
      # test3: whether `axis` arg works as expected.
      result_1d = mi._masked_logmeanexp(
          self.scores,
          self.scores > 0,
          axis=[
              1,
          ],
      )
      self.assertEqual(
          result_1d.shape,
          [
              13,
          ],
      )
  
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 4 for calling the MutualInformationTest obj's underlying member method code with example usage: obj.test_masked_logmeanexp() without return types