| 223 | |
| 224 | |
| 225 | class ApplyRotaryEmbeddingInplace(torch.autograd.Function): |
| 226 | @staticmethod |
| 227 | def forward(x, theta, conj): |
| 228 | _apply_rotary_emb_inplace(x, theta, conj=conj) |
| 229 | return x |
| 230 | |
| 231 | @staticmethod |
| 232 | def setup_context(ctx, inputs, output): |
| 233 | _, theta, conj = inputs |
| 234 | ctx.save_for_backward(theta) |
| 235 | ctx.conj = conj |
| 236 | |
| 237 | @staticmethod |
| 238 | def backward(ctx, grad_output): |
| 239 | theta, = ctx.saved_tensors |
| 240 | _apply_rotary_emb_inplace(grad_output, theta, conj=not ctx.conj) |
| 241 | return grad_output, None, None |
| 242 | |
| 243 | |
| 244 | def apply_rotary_emb_(x, theta): |
nothing calls this directly
no outgoing calls
no test coverage detected