Function _cat_with_padding

gemma_encoder.py:499–516

Concatenate a tensor with a padding tensor of the given value.

(
    tensor: torch.Tensor,
    padding_length: int,
    value: int | float,
)

def _cat_with_padding(
    tensor: torch.Tensor,
    padding_length: int,
    value: int | float,
) -> torch.Tensor:
    """Concatenate a tensor with a padding tensor of the given value."""
    return torch.cat(
        [
            tensor,
            torch.full(
                (1, padding_length),
                value,
                dtype=tensor.dtype,
                device=tensor.device,
            ),
        ],
        dim=1,
    )


def _pad_inputs_for_attention_alignment(model_inputs, pad_token_id, alignment: int = 8):