import torch


def modify_mask_pre_post_softmax(mask):
    """Modify the attention mask to a pre-softmax mask and a post-softmax mask.
    The issue is that when one row out of bsz * num_heads (tgt_max_len, src_max_len) masks is full zeros,
    after softmax the distribution will be all "nan"s, which will cause problem when calculating gradients.

    Thus, we mask out these rows after softmax by multiplying them with 0;
    and we need to modify the pre-softmax as well, since after we get nan, multiplying by 0 is still nan.

    Args:
        mask (torch.Tensor): binary attention mask of size (batch_size * num_heads, tgt_max_len, src_max_len).

    Returns:
        mask_pre_softmax (torch.BoolTensor): binary attention mask of the same size, used pre-softmax.
        mask_post_softmax (torch.FloatTensor): binary attention mask of size (batch_size * num_heads, tgt_max_len),
            to multiple after the softmax.
    """
    mask_post_softmax = mask.new_ones(*mask.size()[:2], 1, dtype=torch.float)
    mask_post_softmax[mask.sum(dim=2) == 0] = 0
    # modify the pre-softmax mask
    mask[(mask.sum(dim=2, keepdim=True) == 0).repeat(1, 1, mask.size(-1))] = 1
    return mask.to(torch.bool), mask_post_softmax