File size: 1,360 Bytes
b66f552 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | import torch
from typing import Tuple
from fla.utils import tensor_cache
@tensor_cache
def prepare_sample_relpos_global_index(
offsets: torch.Tensor
) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]:
lengths = offsets[1:] - offsets[:-1]
S = lengths.numel()
sample_idx_per_token = torch.repeat_interleave(torch.arange(S, device=offsets.device), lengths) # [L]
token_global_idx = torch.arange(offsets[-1], device=offsets.device) # [L]
token_start_idx = offsets[:-1].index_select(0, sample_idx_per_token) # [L]
relpos_in_sample = token_global_idx - token_start_idx
return sample_idx_per_token, relpos_in_sample, token_global_idx, lengths
@tensor_cache
def prepare_sample_relpos_global_index_flat(
offsets: torch.Tensor,
K: int
) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]:
sample_idx_per_token, relpos_in_sample, token_global_idx, lengths = prepare_sample_relpos_global_index(offsets)
sample_idx_flat = sample_idx_per_token[:, None].expand(-1, K).reshape(-1) # [L*K]
relpos_flat = relpos_in_sample[:, None].expand(-1, K).reshape(-1) # [L*K]
global_idx_flat = token_global_idx[:, None].expand(-1, K).reshape(-1) # [L*K]
return sample_idx_flat, relpos_flat, global_idx_flat, lengths
|