def lengths_to_attention_mask(
    lengths: Tensor,
    left_context: Optional[int] = None,
    right_context: Optional[int] = None,
) -> Optional[Tensor]:
    """
    Generate attention mask based on (lengths, left_context, right_context).
    left_context is None means unlimited left context.
    right_context is None means unlimited right context.
    """

    if left_context is None and right_context is None:
        return None

    max_length = int(torch.max(lengths).item())

    # For example, with `max_length` == 5,
    # indices = tensor([
    #     [ 0,  1,  2,  3,  4,  5],
    #     [-1,  0,  1,  2,  3,  4],
    #     [-2, -1,  0,  1,  2,  3],
    #     [-3, -2, -1,  0,  1,  2],
    #     [-4, -3, -2, -1,  0,  1],
    #     [-5, -4, -3, -2, -1,  0],
    # ])

    # In some cases the second torch.arange is created on cpu which causes a
    # failure. Adding the device option to guard against it.
    indices = torch.arange(
        max_length, device=lengths.device, dtype=lengths.dtype
    ).expand(max_length, max_length) - torch.arange(
        max_length, device=lengths.device
    ).view(
        max_length, -1
    )

    # For example, with `max_length` == 5,
    # bool_mask = tensor([
    #     [True, True, True, True, True],
    #     [True, True, True, True, True],
    #     [True, True, True, True, True],
    #     [True, True, True, True, True],
    #     [True, True, True, True, True],
    # ])
    bool_mask = (
        torch.tensor([True]).to(device=lengths.device).expand(max_length, max_length)
    )

    # For example, with `max_length` == 5, left_context == 2
    # left_mask = tensor([
    #     [ True,  True, True, True, True],
    #     [ True,  True, True, True, True],
    #     [ True,  True, True, True, True],
    #     [False,  True, True, True, True],
    #     [False, False, True, True, True],
    # ])
    if left_context is not None:
        left_mask = indices >= -left_context
        bool_mask = bool_mask & left_mask

    # For example, with `max_length` == 5, right_context == 1
    # right_mask = tensor([
    #     [True, True, False, False, False],
    #     [True, True,  True, False, False],
    #     [True, True,  True,  True, False],
    #     [True, True,  True,  True,  True],
    #     [True, True,  True,  True,  True],
    # ])
    if right_context is not None:
        right_mask = indices <= right_context
        bool_mask = bool_mask & right_mask

    bool_mask = (~bool_mask).to(device=lengths.device)
    return bool_mask
Generate attention mask based on (lengths, left_context, right_context).
left_context is None means unlimited left context.
right_context is None means unlimited right context.
fairseq/models/speech_to_text/utils.py