import torch


def masked_softmax(vector: torch.Tensor,
                                     mask: torch.Tensor,
                                     dim: int = -1,
                                     mask_fill_value: float = -1e32) -> torch.Tensor:
    """
    ``torch.nn.functional.softmax(vector)`` does not work if some elements of ``vector`` should be
    masked. This performs a softmax on just the non-masked positions of ``vector``. Passing ``None``
    in for the mask is also acceptable, which is just the regular softmax.
    """
    if mask is None:
        result = torch.softmax(vector, dim=dim)
    else:
        mask = mask.float()
        while mask.dim() < vector.dim():
            mask = mask.unsqueeze(1)
        masked_vector = vector.masked_fill((1 - mask).bool(), mask_fill_value)
        result = torch.softmax(masked_vector, dim=dim)
    return result