import torch


def masked_softmax(x, m=None, dim=-1):
    """
    Softmax with mask.
    :param x: the Tensor to be softmaxed.
    :param m: mask.
    :param dim:
    :return:
    """
    if m is not None:
        m = m.float()
        x = x * m
    e_x = torch.exp(x - torch.max(x, dim=dim, keepdim=True)[0])
    if m is not None:
        e_x = e_x * m
    softmax = e_x / (torch.sum(e_x, dim=dim, keepdim=True) + 1e-6)
    return softmax