| 407 | |
| 408 | |
| 409 | class LinearSpaceAttention(LinearAttention): |
| 410 | def forward(self, x, *args, **kwargs): |
| 411 | x = rearrange(x, "b c ... h w -> b ... h w c") |
| 412 | x, batch_ps = pack_one(x, "* h w c") |
| 413 | x, seq_ps = pack_one(x, "b * c") |
| 414 | |
| 415 | x = super().forward(x, *args, **kwargs) |
| 416 | |
| 417 | x = unpack_one(x, seq_ps, "b * c") |
| 418 | x = unpack_one(x, batch_ps, "* h w c") |
| 419 | return rearrange(x, "b ... h w c -> b c ... h w") |
| 420 | |
| 421 | |
| 422 | class SpaceAttention(Attention): |