import torch


def tsn_sample(num_tokens, num_samples, training):
    """ num_tokens >= num_samples
    args:
        num_tokens: int, num of total tokens
        num_samples: int, num of sampled tokens
        training: bool
  
    returns:
        indexes: tensor, sampled indexes of frames
    """

    if num_samples == 1: return torch.tensor([0], dtype=torch.long)

    base = torch.floor(
        (num_tokens - 1) * torch.arange(0, num_samples).to(torch.float) / (num_samples - 1))

    if training:
        offset_range = base[1:] - base[:-1]
        base[:-1] += torch.rand((offset_range.size()[0],)) * offset_range
        indexes = torch.floor(base)
    else:
        indexes = base
    
    return indexes.to(torch.long)