(x, theta, conj=False)
| 193 | |
| 194 | @flags.compile_wrap |
| 195 | def apply_rotary_emb(x, theta, conj=False): |
| 196 | out_dtype = x.dtype |
| 197 | dtype = reduce(torch.promote_types, (x.dtype, theta.dtype, torch.float32)) |
| 198 | d = theta.shape[-1] |
| 199 | assert d * 2 <= x.shape[-1] |
| 200 | x1, x2, x3 = x[..., :d], x[..., d : d * 2], x[..., d * 2 :] |
| 201 | x1, x2, theta = x1.to(dtype), x2.to(dtype), theta.to(dtype) |
| 202 | cos, sin = torch.cos(theta), torch.sin(theta) |
| 203 | sin = -sin if conj else sin |
| 204 | y1 = x1 * cos - x2 * sin |
| 205 | y2 = x2 * cos + x1 * sin |
| 206 | y1, y2 = y1.to(out_dtype), y2.to(out_dtype) |
| 207 | return torch.cat((y1, y2, x3), dim=-1) |
| 208 | |
| 209 | |
| 210 | @flags.compile_wrap |
nothing calls this directly
no outgoing calls
no test coverage detected