Apply rotary position embeddings.
(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor)
| 61 | |
| 62 | # Rotary Position Embedding |
| 63 | def rotary_embedding_ref(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| 64 | """Apply rotary position embeddings.""" |
| 65 | x1 = x[..., ::2] |
| 66 | x2 = x[..., 1::2] |
| 67 | rx1 = x1 * cos - x2 * sin |
| 68 | rx2 = x1 * sin + x2 * cos |
| 69 | return torch.stack([rx1, rx2], dim=-1).flatten(-2) |
| 70 | |
| 71 | # Parallel Reductions |
| 72 | def reduce_sum_ref(x: torch.Tensor, dim: int = -1) -> torch.Tensor: |
nothing calls this directly
no outgoing calls
no test coverage detected