import torch
from typing import Tuple


def get_sizes_and_size_with_padding(mask: torch.Tensor) -> Tuple[torch.Tensor, int]:
    """

    From a batch of tensors indicating padded positions (where the tensor is equal to 0)
    determine the sizes and sizes with padding to be used to generate masks.

    Example with tensor [[1, 1, 1, 0], [1, 1, 0, 0]]:
        -> ([3, 2], 4)
        The first tensor has an unpadded length of 3, the second tensor has an unpadded length of 2.
        Both tensors are padded to length 4.

    :param mask: tensor of shape [B, S]
    :return: tensor indicating sizes of the tensors in the batch, padded size
    """
    assert mask.dim() == 2, f"expected mask to have two dimensions, but got a shape of {mask.shape}"
    B, S = mask.size()
    return mask.sum(dim=1).long(), S