(x)
| 145 | |
| 146 | |
| 147 | def rotate_half(x): |
| 148 | x = rearrange(x, "... (d r) -> ... d r", r=2) |
| 149 | x1, x2 = x.unbind(dim=-1) |
| 150 | x = torch.stack((-x2, x1), dim=-1) |
| 151 | return rearrange(x, "... d r -> ... (d r)") |
| 152 | |
| 153 | |
| 154 | class RotaryPositionalEmbedding1D(nn.Module): |