import torch
from typing import Optional


def masked_softmax(
    x: torch.Tensor, mask: torch.Tensor, dim: Optional[int] = 1
) -> torch.Tensor:
    """Compute softmax of a tensor using a mask.

    Parameters
    ----------
    x
        Argument of softmax.
    mask
        Mask tensor.
    dim
        Dimension along which softmax will be computed.

    Returns
    -------
    torch.Tensor
        Masked softmax of `x`.
    """
    X = (mask * x).masked_fill(mask == 0, -float("inf"))
    return torch.softmax(X, dim)