def mask_joint_logits(input_mask, start_end_logits):
  """Masks logits based on input mask and valid start/end combinations."""
  _, _, length = modeling.get_shape_list(start_end_logits, expected_rank=3)

  mask = tf.TensorArray(input_mask.dtype, size=length, dynamic_size=False)
  for i in range(length):
    mask = mask.write(i, input_mask)
    # The permitted span length is determined by the existing mask combined
    # with its being shifted up by one.
    input_mask = input_mask * tf.pad(input_mask[:, 1:], [[0, 0], [0, 1]])
  mask = mask.stack()
  mask = tf.transpose(mask, [1, 2, 0])
  mask.shape.assert_is_compatible_with(start_end_logits.shape)

  start_end_logits -= 1e6 * tf.cast(1 - mask, tf.float32)
  return start_end_logits
Masks logits based on input mask and valid start/end combinations.
language/capwap/synthetic/generate_answers.py