code: def _check_mask_shape(logu, joint_sample_mask):
      """Validate the mask specifying joint samples of the score matrix."""
      row_sum = tf.reduce_sum(
          tf.cast(joint_sample_mask, logu.dtype), axis=[-1]
      )
      assertion_list = []
      assertion_list.append(
          assert_util.assert_equal(
              logu.shape,
              joint_sample_mask.shape,
              message="logu should be of the same shape as joint_sample_mask.",
          )
      )
      assertion_list.append(
          assert_util.assert_equal(
              row_sum,
              tf.ones_like(row_sum),
              message="joint_sample_mask should be one-hot for each row.",
          )
      )
      return assertion_list
  
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 1 for calling a global function with example usage: _check_mask_shape(logu, joint_sample_mask) and returns: assertion_list