import torch


def stable_softmax(scores, mask=None, epsilon=1e-9):
    """
    :param scores: tensor of shape (batch_size, sequence_length)
    :param mask: (optional) binary tensor in case of padded sequences, shape: (batch_size, sequence_length)
    :param epsilon: epsilon to be added to the normalization factor
    :return: probability tensor of shape (batch_size, sequence_length)
    """
    batch, seq = scores.size()

    # Numerically stable masked-softmax
    maxvec, _ = scores.max(dim=1, keepdim=True)
    scores = torch.exp(scores - maxvec)  # greatest exp value is up to 1.0 to avoid overflow

    if mask is not None:
        scores = scores * mask.view(batch, seq).float()

    sums = torch.sum(scores, dim=1, keepdim=True) + epsilon  # add epsilon to avoid div by zero in case of underflow
    prob = scores / sums

    return prob