(window_size, x)
| 265 | # Shifted window attention |
| 266 | |
| 267 | def window(window_size, x): |
| 268 | *b, h, w, c = x.shape |
| 269 | x = torch.reshape( |
| 270 | x, |
| 271 | (*b, h // window_size, window_size, w // window_size, window_size, c), |
| 272 | ) |
| 273 | x = torch.permute( |
| 274 | x, |
| 275 | (*range(len(b)), -5, -3, -4, -2, -1), |
| 276 | ) |
| 277 | return x |
| 278 | |
| 279 | |
| 280 | def unwindow(x): |