code: def _check_and_get_mask(
      logu, joint_sample_mask=None, validate_args=False
  ):
      """Helper function for creating masks for joint/marginal samples.
  
      The function is designed to do:
        - when `joint_sample_mask` is provided, check and validate the mask.
        - when `joint_sample_mask` is not provided, generate a default mask.
  
      Variational bounds on mutual information require knowing which
      elements of the score matrix correspond to positive elements
      (sampled from joint distribution `p(x,y)`) and negative elements
      (sampled from marginal distribution `p(x)p(y)`). By default, we assume
      that the diagonal elements of scores contain positaive pairs, and all
      other elements are negatives.
  
      Args:
        logu: `float`-like `Tensor` representing scores to be masked.
        joint_sample_mask: `bool`-like `Tensor` of the same shape of `logu`
          masking the joint samples by `True`, i.e. samples from joint
          distributions `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.
  
      Returns:
        logu: `float`-like `Tensor` based on input `logu`.
        joint_sample_mask: `bool`-like `Tensor` for joint samples.
      """
      with tf.name_scope("get_default_mask"):
          logu = tf.convert_to_tensor(
              logu, dtype_hint=tf.float32, name="logu"
          )
  
          if joint_sample_mask is None:
              num_rows, num_cols = tf.unstack(tf.shape(logu)[-2:])
              joint_sample_mask = tf.eye(
                  num_rows, num_cols, dtype=tf.bool
              )
          else:
              joint_sample_mask = tf.convert_to_tensor(
                  joint_sample_mask,
                  dtype_hint=tf.bool,
                  name="joint_sample_mask",
              )
              with tf.control_dependencies(
                  _check_mask_shape(logu, joint_sample_mask)
                  if validate_args
                  else []
              ):
                  joint_sample_mask = tf.identity(joint_sample_mask)
          return logu, joint_sample_mask
  
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 2 for calling a global function with example usage: _check_and_get_mask(logu, joint_sample_mask, validate_args) and returns: logu,  joint_sample_mask