(n_h_w, n_w_w, w_h, w_w, shift, device=None)
| 298 | |
| 299 | @lru_cache |
| 300 | def make_shifted_window_masks(n_h_w, n_w_w, w_h, w_w, shift, device=None): |
| 301 | ph_coords = torch.arange(n_h_w, device=device) |
| 302 | pw_coords = torch.arange(n_w_w, device=device) |
| 303 | h_coords = torch.arange(w_h, device=device) |
| 304 | w_coords = torch.arange(w_w, device=device) |
| 305 | patch_h, patch_w, q_h, q_w, k_h, k_w = torch.meshgrid( |
| 306 | ph_coords, |
| 307 | pw_coords, |
| 308 | h_coords, |
| 309 | w_coords, |
| 310 | h_coords, |
| 311 | w_coords, |
| 312 | indexing="ij", |
| 313 | ) |
| 314 | is_top_patch = patch_h == 0 |
| 315 | is_left_patch = patch_w == 0 |
| 316 | q_above_shift = q_h < shift |
| 317 | k_above_shift = k_h < shift |
| 318 | q_left_of_shift = q_w < shift |
| 319 | k_left_of_shift = k_w < shift |
| 320 | m_corner = ( |
| 321 | is_left_patch |
| 322 | & is_top_patch |
| 323 | & (q_left_of_shift == k_left_of_shift) |
| 324 | & (q_above_shift == k_above_shift) |
| 325 | ) |
| 326 | m_left = is_left_patch & ~is_top_patch & (q_left_of_shift == k_left_of_shift) |
| 327 | m_top = ~is_left_patch & is_top_patch & (q_above_shift == k_above_shift) |
| 328 | m_rest = ~is_left_patch & ~is_top_patch |
| 329 | m = m_corner | m_left | m_top | m_rest |
| 330 | return m |
| 331 | |
| 332 | |
| 333 | def apply_window_attention(window_size, window_shift, q, k, v, scale=None): |
no outgoing calls
no test coverage detected