def pad_sequence(
    sequence: Tensor,
    time_axis: int,
    extra_left_context: int = 0,
    extra_right_context: int = 0,
) -> Tensor:
    """Pad extra left/right contexts to the sequence."""

    if extra_left_context == 0 and extra_right_context == 0:
        return sequence

    tensors_to_concat = []

    if extra_left_context:
        size = (extra_left_context,)
        fill_value = 0
        indices = torch.full(
            size=size,
            fill_value=fill_value,
            dtype=torch.long,
            device=sequence.device,
        )
        left_padding = torch.index_select(sequence, time_axis, indices)
        tensors_to_concat.append(left_padding)

    tensors_to_concat.append(sequence)

    # NOTE(cfyeh): for efficiency reason we pad 0 instead of the last frame for
    #              extra right contexts.
    if extra_right_context:
        size = list(sequence.shape)
        size[time_axis] = extra_right_context
        right_padding = torch.zeros(size, dtype=sequence.dtype, device=sequence.device)
        tensors_to_concat.append(right_padding)

    padded_sequence = torch.cat(tensors_to_concat, dim=time_axis)
    return padded_sequence
Pad extra left/right contexts to the sequence.
understand_icl/fairseq/fairseq/models/speech_to_text/utils.py