| 255 | |
| 256 | |
| 257 | class TokenShift(Module): |
| 258 | @beartype |
| 259 | def __init__(self, fn: Module): |
| 260 | super().__init__() |
| 261 | self.fn = fn |
| 262 | |
| 263 | def forward(self, x, **kwargs): |
| 264 | x, x_shift = x.chunk(2, dim=1) |
| 265 | x_shift = pad_at_dim(x_shift, (1, -1), dim=2) # shift time dimension |
| 266 | x = torch.cat((x, x_shift), dim=1) |
| 267 | return self.fn(x, **kwargs) |
| 268 | |
| 269 | |
| 270 | # rmsnorm |