| 19 | |
| 20 | @flags.compile_wrap |
| 21 | def apply_rotary_emb(freqs, t, start_index=0, scale=1.0): |
| 22 | freqs = freqs.to(t) |
| 23 | rot_dim = freqs.shape[-1] |
| 24 | end_index = start_index + rot_dim |
| 25 | assert rot_dim <= t.shape[-1], f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}" |
| 26 | t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:] |
| 27 | t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale) |
| 28 | return torch.cat((t_left, t, t_right), dim=-1) |
| 29 | |
| 30 | |
| 31 | def centers(start, stop, num, dtype=None, device=None): |