import torch
from typing import Optional


def num2mask(
        nums: torch.Tensor,
        max_length: Optional[int] = None
) -> torch.Tensor:
    """
    E.g. input a tensor [2, 3, 4], return [[T T F F], [T T T F], [T T T T]]
    :param nums: Shape [batch]
    :param max_length: maximum length. if not provided, will choose the largest number from nums.
    :return: 2D binary mask.
    """
    shape_backup = nums.shape
    nums = nums.flatten()
    max_length = max_length or int(nums.max())
    batch_size = len(nums)
    range_nums = torch.arange(0, max_length, device=nums.device).unsqueeze(0).expand([batch_size, max_length])
    ret = (range_nums.T < nums).T
    return ret.reshape(*shape_backup, max_length)