def mask_attention(attention, seq_len1, seq_len2):
  """Masks an attention matrix.

  Args:
    attention: <tf.float32>[batch, seq_len1, seq_len2]
    seq_len1: <tf.int32>[batch]
    seq_len2: <tf.int32>[batch]

  Returns:
    the masked scores <tf.float32>[batch, seq_len1, seq_len2]
  """
  dim1 = tensor_utils.shape(attention, 1)
  dim2 = tensor_utils.shape(attention, 2)
  m1 = tf.sequence_mask(seq_len1, dim1)
  m2 = tf.sequence_mask(seq_len2, dim2)
  joint_mask = tf.logical_and(tf.expand_dims(m1, 2), tf.expand_dims(m2, 1))
  return ops.mask_logits(attention, joint_mask)
Masks an attention matrix.

  Args:
    attention: <tf.float32>[batch, seq_len1, seq_len2]
    seq_len1: <tf.int32>[batch]
    seq_len2: <tf.int32>[batch]

  Returns:
    the masked scores <tf.float32>[batch, seq_len1, seq_len2]
  
language/boolq/run_recurrent_model_boolq.py