def mask_renormalize(probs, mask):
  """Renormalizes probs with a mask so that the unmasked entries sum to 1.

  Args:
    probs (torch.FloatTensor): batch of probability distributions of shape
      (batch_dim1, batch_dim2, ..., num_elems).
    mask (torch.ByteTensor): mask of same shape as probs, where entries = 0
      should be masked out.

  Returns:
    renormalized_probs (torch.FloatTensor): tensor of same shape as probs. Each
      batch row (last dim) sums to 1, where masked entries have 0 prob. If all
      entries in a batch are masked, the batch row sums to 0.
  """
  masked_probs = probs * mask.float()
  renormalized = masked_probs / (masked_probs.sum(-1, keepdim=True) + 1e-8)
  return renormalized