import torch


def create_mask_from_length(length_tensor, mask_size, zeros_at_end=True):
    """
    Creates a binary mask based on length.

    :param length_tensor: ND Tensor containing the lengths.
    :param mask_size: Int specifying the mask size. Usually the largest length.
    :param zeros_at_end: Whether to put the zeros of the mask at the end.
    :return: (N+1)D Int Tensor (..., mask_size).
    """

    if zeros_at_end:
        mask = torch.arange(0, mask_size, dtype=torch.int, device=length_tensor.device)
    else:
        mask = torch.arange(mask_size - 1, -1, step=-1, dtype=torch.int, device=length_tensor.device)

    mask = mask.int().view([1] * (len(length_tensor.shape)) + [-1])

    return mask < length_tensor.int().unsqueeze(-1)