(x, patch_size)
| 197 | |
| 198 | |
| 199 | def patchify(x, patch_size): |
| 200 | if patch_size == 1: |
| 201 | return x |
| 202 | if x.dim() == 4: |
| 203 | x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) |
| 204 | elif x.dim() == 5: |
| 205 | x = rearrange(x, |
| 206 | "b c f (h q) (w r) -> b (c r q) f h w", |
| 207 | q=patch_size, |
| 208 | r=patch_size) |
| 209 | else: |
| 210 | raise ValueError(f"Invalid input shape: {x.shape}") |
| 211 | return x |
| 212 | |
| 213 | |
| 214 | def unpatchify(x, patch_size): |