xma.functional.sequence_packing

pack_sequence(inputs: Sequence[Tensor], cu_seqlens: Tensor, total_tokens: int, padding_side: str = 'left', *, kernel_backend: KernelBackend | None = None) → Sequence[Tensor]

pack tensors

Parameters:

• inputs (Sequence[torch.Tensor]) – list of tensors

• cu_seqlens (torch.Tensor) – cumulative sequence length

• total_tokens (int) – total number of tokens

• padding_side (str) – padding side

• kernel_backend (KernelBackend | None) – KernelBackend

Returns:

list of packed tensors

Return type:

Sequence[Tensor]

unpack_sequence(inputs: Sequence[Tensor], cu_seqlens: Tensor, batch_size: int, sequence_length: int, padding_side: str = 'left', *, kernel_backend: KernelBackend | None = None) → Sequence[Tensor]

unpack tensors

Parameters:

• inputs (Sequence[torch.Tensor]) – list of tensors

• cu_seqlens (torch.Tensor) – cumulative sequence length

• batch_size (int) – batch size

• sequence_length (int) – sequence length

• padding_side (str) – padding side

• kernel_backend (KernelBackend | None) – KernelBackend

Returns:

list of unpacked tensors

Return type:

Sequence[Tensor]