| 94 | """Partition / reverse-partition helpers for 5-D tensors (B,F,H,W,C).""" |
| 95 | @staticmethod |
| 96 | def partition(x: torch.Tensor, win: Tuple[int, int, int]): |
| 97 | B, F, H, W, C = x.shape |
| 98 | wf, wh, ww = win |
| 99 | assert F % wf == 0 and H % wh == 0 and W % ww == 0, "Dims must divide by window size." |
| 100 | x = x.view(B, F // wf, wf, H // wh, wh, W // ww, ww, C) |
| 101 | x = x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous() |
| 102 | return x.view(-1, wf * wh * ww, C) |
| 103 | |
| 104 | @staticmethod |
| 105 | def reverse(windows: torch.Tensor, win: Tuple[int, int, int], orig: Tuple[int, int, int]): |