import torch


def sequence_mask_mxlen(lengths: torch.Tensor, max_len: int) -> torch.Tensor:
    """Generate a sequence mask of shape `BxT` based on the given lengths, with a maximum value

    This function primarily exists to make ONNX tracing work better
    :param lengths: A `B` tensor containing the lengths of each example
    :param max_len: The maximum width (length) allowed in this mask (default to None)
    :return: A mask
    """
    lens = lengths.cpu()
    max_len_v = max_len
    # 1 x T
    row = torch.arange(0, max_len_v).type_as(lens).view(1, -1)
    # B x 1
    col = lens.view(-1, 1)
    # Broadcast to B x T, compares increasing number to max
    mask = row < col
    return mask