| import torch |
|
|
| def move_masked_to_left_brute_force(tensor, mask): |
| results = [] |
| new_mask = [] |
| for i in range(tensor.shape[0]): |
| t = torch.cat([tensor[i][mask[i]], torch.zeros_like(tensor[i][~mask[i]])]) |
| results.append(t) |
| l = [True] * mask[i].sum() + [False] * (~mask[i]).sum() |
| new_mask.append(l) |
| results = torch.stack(results) |
| new_mask = torch.tensor(new_mask, dtype=torch.bool) |
| max_length = mask.sum(dim=1).max() |
|
|
| results = results[:, :max_length] |
| new_mask = new_mask[:, :max_length] |
| return results, new_mask |
|
|
|
|
| def move_masked_to_left_ids(tensor, mask, pad_zero=True): |
| masked_index = mask.cumsum(dim=1) - 1 |
| unmasked_index = (~mask).cumsum(dim=1) - 1 |
| unmasked_index += mask.sum(dim=1).unsqueeze(1) |
| s2t_index = torch.where(mask, masked_index, unmasked_index) |
| t2s_index = torch.argsort(s2t_index, dim=1) |
| result = torch.gather(tensor, 1, t2s_index) |
| |
| length = mask.sum(dim=1) |
| result = result[:, :length.max()] |
| new_mask = torch.arange(result.shape[1], device=length.device).unsqueeze(0) < length.unsqueeze(1) |
| if pad_zero: |
| result[~new_mask] = 0 |
| return result, new_mask |
|
|
|
|
| def move_masked_to_left(tensor, mask, pad_zero=True): |
| masked_index = mask.cumsum(dim=1) - 1 |
| unmasked_index = (~mask).cumsum(dim=1) - 1 |
| unmasked_index += mask.sum(dim=1).unsqueeze(1) |
| s2t_index = torch.where(mask, masked_index, unmasked_index) |
| t2s_index = torch.argsort(s2t_index, dim=1) |
| result = torch.gather(tensor, 1, t2s_index.unsqueeze(2).expand(-1, -1, tensor.shape[2])) |
| |
| length = mask.sum(dim=1) |
| result = result[:, :length.max()] |
| new_mask = torch.arange(result.shape[1], device=length.device).unsqueeze(0) < length.unsqueeze(1) |
| if pad_zero: |
| result[~new_mask] = 0 |
| return result, new_mask |
|
|
| def get_mask_of_last_masked_index_brute_force(mask, length): |
| results = [] |
| for i in range(mask.shape[0]): |
| len = length if isinstance(length, int) else length[i].item() |
| l = [False] * mask.shape[1] |
| for j in range(mask.shape[1] - 1, -1, -1): |
| if mask[i][j] and len > 0: |
| l[j] = True |
| len -= 1 |
| else: |
| l[j] = False |
| results.append(l) |
| return torch.tensor(results, dtype=torch.bool) |
|
|
| def get_mask_of_last_masked_index(mask, length): |
| cumsum = mask.cumsum(dim=1) |
| new_length = mask.sum(dim=1) - length |
| last_masked_index = (cumsum > new_length.unsqueeze(1)) & mask |
| return last_masked_index |
|
|
|
|
| def test_move_masked_to_left(): |
| b = 10 |
| n = 20 |
| tensor = torch.randn(b, n, 5) |
| mask = torch.randint(0, 2, (b, n)).bool() |
| |
| |
| result_1, mask_1 = move_masked_to_left(tensor, mask) |
| result_2, mask_2 = move_masked_to_left_brute_force(tensor, mask) |
| |
| assert (result_1 == result_2).all() |
| assert (mask_1 == mask_2).all() |
| assert (mask.sum(dim=1) == mask_1.sum(dim=1)).all() |
|
|
| for i in range(b): |
| l = mask[i].sum() |
| assert (mask_1[i][:l].all() == True) |
| assert (mask_1[i][l:].any() == False) |
|
|
| def test_get_mask_of_last_masked_index(): |
| b = 10 |
| n = 20 |
| mask = torch.randint(0, 2, (b, n)).bool() |
| length = torch.randint(0, n//2, (b,)) |
| last_masked_index_1 = get_mask_of_last_masked_index(mask, length) |
| last_masked_index_2 = get_mask_of_last_masked_index_brute_force(mask, length) |
| assert (last_masked_index_1 == last_masked_index_2).all() |
|
|
| if __name__ == '__main__': |
| test_move_masked_to_left() |
| test_get_mask_of_last_masked_index() |
