| 48 | |
| 49 | |
| 50 | def create_mask(window_size, displacement, upper_lower, left_right): |
| 51 | mask = torch.zeros(window_size ** 2, window_size ** 2) |
| 52 | |
| 53 | if upper_lower: |
| 54 | mask[-displacement * window_size:, :-displacement * window_size] = float('-inf') |
| 55 | mask[:-displacement * window_size, -displacement * window_size:] = float('-inf') |
| 56 | |
| 57 | if left_right: |
| 58 | mask = rearrange(mask, '(h1 w1) (h2 w2) -> h1 w1 h2 w2', h1=window_size, h2=window_size) |
| 59 | mask[:, -displacement:, :, :-displacement] = float('-inf') |
| 60 | mask[:, :-displacement, :, -displacement:] = float('-inf') |
| 61 | mask = rearrange(mask, 'h1 w1 h2 w2 -> (h1 w1) (h2 w2)') |
| 62 | |
| 63 | return mask |
| 64 | |
| 65 | |
| 66 | def get_relative_distances(window_size): |