code: def mask_matrix(x, is_missing=None):
      """Copies a matrix, replacing masked-out rows/cols from the identity matrix.
  
      Args:
        x: A Tensor of shape `[..., n, n]`, representing a batch of n-by-n matrices.
        is_missing: A boolean Tensor of shape `[..., n]`, representing a batch of
          masks. If `is_missing` is None, `x` is returned.
      Returns:
        A Tensor of shape `[..., n, n]`, representing a batch of n-by-n matrices.
        For each batch member `r`, element `r[i, j]` equals `eye(n)[i, j]` if
        dimension `i` or `j` is True in the corresponding input mask.  Otherwise,
        `r[i, j]` equals the corresponding element from `x`.
      """
      if is_missing is None:
          return x
  
      x = tf.convert_to_tensor(x)
      is_missing = tf.convert_to_tensor(is_missing, dtype=tf.bool)
  
      n = ps.dimension_size(x, -1)
  
      return tf.where(
          is_missing[..., tf.newaxis] | is_missing[..., tf.newaxis, :],
          tf.eye(n, dtype=x.dtype),
          x,
      )
  
file_path: probability/tensorflow_probability/python/math/psd_kernels/internal/util.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 7 for calling a global function with example usage: mask_matrix(x, is_missing) and returns: tf, x