def METHOD_NAME(reserved_length: int,
                       seq_length: int,
                       batch_size: int,
                       layers: int = 1,
                       device: str = 'cpu',
                       attn_mask: torch.Tensor = None):
    assert attn_mask is not None
    prob_dist = torch.ones((layers * batch_size, seq_length), device=device)
    sampled_indices = torch.multinomial(prob_dist, reserved_length)
    sampled_indices = sampled_indices.reshape(layers, batch_size, reserved_length).to(torch.int32)
    global random_ltd_module
    if random_ltd_module is None:
        random_ltd_module = RandomLTDBuilder().load()
    sampled_indices = random_ltd_module.token_sort_(sampled_indices, seq_length)
    dtype = sampled_indices.dtype
    sampled_indices = sampled_indices.to(torch.long)
    new_mask = []
    for l in range(layers):
        tmp_mask_list = []
        for i in range(batch_size):
            mask_tmp = attn_mask[i:i + 1, :, sampled_indices[l][i], :]
            tmp_mask_list.append(mask_tmp[:, :, :, sampled_indices[l][i]])
        new_mask.append(torch.cat(tmp_mask_list, dim=0))
    return sampled_indices.to(dtype), new_mask