(self, x, *args, **kwargs)
| 421 | |
| 422 | class SpaceAttention(Attention): |
| 423 | def forward(self, x, *args, **kwargs): |
| 424 | x = rearrange(x, "b c t h w -> b t h w c") |
| 425 | x, batch_ps = pack_one(x, "* h w c") |
| 426 | x, seq_ps = pack_one(x, "b * c") |
| 427 | |
| 428 | x = super().forward(x, *args, **kwargs) |
| 429 | |
| 430 | x = unpack_one(x, seq_ps, "b * c") |
| 431 | x = unpack_one(x, batch_ps, "* h w c") |
| 432 | return rearrange(x, "b t h w c -> b c t h w") |
| 433 | |
| 434 | |
| 435 | class TimeAttention(Attention): |
nothing calls this directly
no test coverage detected