code: def lower_bound_info_nce(
      logu, joint_sample_mask=None, validate_args=False, name=None
  ):
      """InfoNCE lower bound on mutual information.
  
      InfoNCE lower bound is proposed in [van den Oord et al. (2018)][1]
      based on noise contrastive estimation (NCE).
      ```none
      I(X; Y) >= 1/K sum(i=1:K, log( p_joint[i] / p_marginal[i])),
      ```
      where the numerator and the denominator are, respectively,
      ```none
      p_joint[i] = p(x[i] | y[i]) = exp( f(x[i], y[i]) ),
      p_marginal[i] = 1/K sum(j=1:K, p(x[i] | y[j]) )
                    = 1/K sum(j=1:K, exp( f(x[i], y[j]) ) ),
      ```
      and `(x[i], y[i]), i=1:K` are samples from joint distribution `p(x, y)`.
      Pairs of points (x, y) are scored using a critic function `f`.
  
      Example:
  
      `X`, `Y` are samples from a joint Gaussian distribution, with
      correlation `0.8` and both of dimension `1`.
  
      ```python
      batch_size, rho, dim = 10000, 0.8, 1
      y, eps = tf.split(
          value=tf.random.normal(shape=(2 * batch_size, dim), seed=7),
          num_or_size_splits=2, axis=0)
      mean, conditional_stddev = rho * y, tf.sqrt(1. - tf.square(rho))
      x = mean + conditional_stddev * eps
  
      # Conditional distribution of p(x|y)
      conditional_dist = tfd.MultivariateNormalDiag(
          mean, scale_diag=conditional_stddev * tf.ones((batch_size, dim)))
  
      # Scores/unnormalized likelihood of pairs of samples `x[i], y[j]`
      # (The scores has its shape [x_batch_size, distibution_batch_size]
      # as the `lower_bound_info_nce` requires `scores[i, j] = f(x[i], y[j])
      # = log p(x[i] | y[j])`.)
      scores = conditional_dist.log_prob(x[:, tf.newaxis, :])
  
      # Mask for joint samples
      joint_sample_mask = tf.eye(batch_size, dtype=bool)
  
      # InfoNCE lower bound on mutual information
      lower_bound_info_nce(logu=scores, joint_sample_mask=joint_sample_mask)
      ```
  
      Args:
        logu: `float`-like `Tensor` of size `[batch_size_1, batch_size_2]`
          representing critic scores (scores) for pairs of points (x, y) with
          `logu[i, j] = f(x[i], y[j])`.
        joint_sample_mask: `bool`-like `Tensor` of the same size as `logu`
          masking the positive samples by `True`, i.e. samples from joint
          distribution `p(x, y)`.
          Default value: `None`. By default, an identity matrix is constructed as
          the mask.
        validate_args: Python `bool`, default `False`. Whether to validate input
          with asserts. If `validate_args` is `False`, and the inputs are invalid,
          correct behavior is not guaranteed.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., 'lower_bound_info_nce').
  
      Returns:
        lower_bound: `float`-like `scalar` for lower bound on mutual information.
  
      #### References
  
      [1]: Aaron van den Oord, Yazhe Li, Oriol Vinyals. Representation
           Learning with Contrastive Predictive Coding. _arXiv preprint
           arXiv:1807.03748_, 2018. https://arxiv.org/abs/1807.03748.
      """
  
      with tf.name_scope(name or "lower_bound_info_nce"):
          # Follow the notation of eq.(12) of Poole et al. (2019)
          # On Variational Bounds of Mutual Information,
          # https://arxiv.org/abs/1905.06922, where the expectation is taken by
          # sampling.
          # The first term is `1/K * sum(i=1:K, f(x[i], y[i])`, where `K` is the
          # `batch_size` and `(x[i], y[i])` is the joint sample.
          # The second term is `1/K * sum(i=1:K, log(sum(j=1:K, exp(f(x[i], y[j]))))`,
          # where the joint samples are when `i=j`, and the marginal ones are `i!=j`.
  
          with tf.control_dependencies(
              _maybe_assert_float_matrix(logu, validate_args)
          ):
              if joint_sample_mask is None:
                  logu = tf.convert_to_tensor(
                      logu, dtype_hint=tf.float32, name="logu"
                  )
                  joint_term = tf.reduce_mean(
                      tf.linalg.diag_part(logu), axis=[-1]
                  )
  
              else:
                  logu, joint_sample_mask = _check_and_get_mask(
                      logu,
                      joint_sample_mask,
                      validate_args=validate_args,
                  )
                  joint_term = tf.reduce_mean(
                      tf.boolean_mask(logu, joint_sample_mask),
                      axis=[-1],
                  )
  
              log_n = tf.math.log(tf.cast(logu.shape[-1], logu.dtype))
              marginal_term = (
                  tf.reduce_mean(
                      tf.reduce_logsumexp(logu, axis=[-1]), axis=[-1]
                  )
                  - log_n
              )
              return joint_term - marginal_term
  
file_path: probability/tensorflow_probability/python/vi/mutual_information.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 7 for calling a global function with example usage: lower_bound_info_nce(logu, joint_sample_mask, validate_args, name) and returns: unknown