import torch


def masked_softmax(logits: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """
    Computes the masked softmax from logits.

    :param logits: (N x ... x K) tensor
    :param mask: (N x ... x K) tensor
    :return: (N x ... x K) tensor
    """
    maxima, _ = torch.max(logits, dim=-1, keepdim=True)
    exps = torch.exp(logits - maxima)

    # Note: in-place operator must not be used here!
    exps = exps * mask
    return exps / torch.sum(exps, keepdim=True, dim=-1)