import torch


def make_pad_mask(lengths: torch.Tensor, le : bool = True) -> torch.Tensor:
    """Make mask tensor containing indices of padded part.

    See description of make_non_pad_mask.

    Args:
        lengths (torch.Tensor): Batch of lengths (B,).
    Returns:
        torch.Tensor: Mask tensor containing indices of padded part.

    Examples:
        >>> lengths = [5, 3, 2]
        >>> make_pad_mask(lengths)
        # masks = [[0, 0, 0, 0 ,0],
        #          [0, 0, 0, 1, 1],
        #          [0, 0, 1, 1, 1]]
        masks = [[1, 1, 1, 1, 1],
                 [1, 1, 1, 0, 0],
                 [1, 1, 0, 0, 0]]
    """
    batch_size = int(lengths.size(0))
    max_len = int(lengths.max().item())
    seq_range = torch.arange(0,
                             max_len,
                             dtype=torch.int64,
                             device=lengths.device)
    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
    seq_length_expand = lengths.unsqueeze(-1)
    # mask = seq_range_expand >= seq_length_expand
    # fix: torch.float32 -> torch.int32
    if le:
        mask = (seq_range_expand < seq_length_expand).type(torch.int32)
    else:
        mask = (seq_range_expand >= seq_length_expand).type(torch.int32)

    # print(mask)
    return mask