(x, patch_size)
| 283 | |
| 284 | |
| 285 | def patchify(x, patch_size): |
| 286 | if patch_size == 1: |
| 287 | return x |
| 288 | if x.dim() == 4: |
| 289 | x = rearrange( |
| 290 | x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) |
| 291 | elif x.dim() == 5: |
| 292 | x = rearrange( |
| 293 | x, |
| 294 | "b c f (h q) (w r) -> b (c r q) f h w", |
| 295 | q=patch_size, |
| 296 | r=patch_size, |
| 297 | ) |
| 298 | else: |
| 299 | raise ValueError(f"Invalid input shape: {x.shape}") |
| 300 | |
| 301 | return x |
| 302 | |
| 303 | |
| 304 | def unpatchify(x, patch_size): |