def mask_logits(vec, mask):
  """Mask `vec` in log-space.

  Elements in `vec` that are not in `mask` will be set to be very
  negative, so that when no longer in log-space (e.i., after `tf.exp(vec)`)
  their values will be very close to zero.

  Args:
    vec: <float32>[...] tensor to mask
    mask: Either a None (in which case this is a no-op), a boolean or 0/1 float
      mask that matches all, or all but the last, dimensions of `vec`, or 1-D
      integer length mask

  Raises:
    ValueError: If `mask` cannot be matched to `vec`
  Returns:
    masked: vec:<float32>[...]
  """
  if mask is None:
    return vec

  if mask.dtype == tf.int32:
    # Assume `mask` holds sequence lengths
    if len(vec.shape) not in [2, 3]:
      raise ValueError("Can't use a length mask on tensor of rank>3")
    mask = tf.sequence_mask(mask, tf.shape(vec)[1], tf.float32)
  else:
    mask = tf.to_float(mask)

  if len(mask.shape) == (len(vec.shape) - 1):
    mask = tf.expand_dims(mask, len(vec.shape) - 1)

  return vec * mask - (1 - mask) * 1E20
Mask `vec` in log-space.

  Elements in `vec` that are not in `mask` will be set to be very
  negative, so that when no longer in log-space (e.i., after `tf.exp(vec)`)
  their values will be very close to zero.

  Args:
    vec: <float32>[...] tensor to mask
    mask: Either a None (in which case this is a no-op), a boolean or 0/1 float
      mask that matches all, or all but the last, dimensions of `vec`, or 1-D
      integer length mask

  Raises:
    ValueError: If `mask` cannot be matched to `vec`
  Returns:
    masked: vec:<float32>[...]
  
language/boolq/utils/ops.py