Rotary position embedding: apply cos/sin rotation to tensor. `cos` and `sin` have shape `(seq_len, head_dim/2)` or compatible broadcast shape.
(&self, x: &Tensor, cos: &Tensor, sin: &Tensor)
| 442 | /// Rotary position embedding: apply cos/sin rotation to tensor. |
| 443 | /// `cos` and `sin` have shape `(seq_len, head_dim/2)` or compatible broadcast shape. |
| 444 | fn rope(&self, x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result<Tensor> { |
| 445 | // Fast path: F32 CPU data — process pairs directly |
| 446 | if x.dtype() == DType::F32 && cos.dtype() == DType::F32 { |
| 447 | let x_shape = x.dims(); |
| 448 | let head_dim = *x_shape.last().unwrap_or(&0); |
| 449 | let half = head_dim / 2; |
| 450 | let x_data = x.contiguous()?.flatten_all()?.to_vec1::<f32>()?; |
| 451 | let cos_data = cos.contiguous()?.flatten_all()?.to_vec1::<f32>()?; |
| 452 | let sin_data = sin.contiguous()?.flatten_all()?.to_vec1::<f32>()?; |
| 453 | let total_vecs = x_data.len() / head_dim; |
| 454 | let cos_stride = cos_data.len() / half; // number of sequence positions in cos |
| 455 | let mut out = vec![0f32; x_data.len()]; |
| 456 | for v in 0..total_vecs { |
| 457 | let x_off = v * head_dim; |
| 458 | // Determine which cos/sin row to use (seq position) |
| 459 | let seq_idx = if cos_stride > 1 { |
| 460 | // x is (batch, heads, seq, dim), cos is (seq, half) |
| 461 | let seq_len = if x_shape.len() >= 3 { |
| 462 | x_shape[x_shape.len() - 2] |
| 463 | } else { |
| 464 | 1 |
| 465 | }; |
| 466 | (v % seq_len) * half |
| 467 | } else { |
| 468 | 0 |
| 469 | }; |
| 470 | for i in 0..half { |
| 471 | let c = cos_data[seq_idx + i]; |
| 472 | let s = sin_data[seq_idx + i]; |
| 473 | let x1 = x_data[x_off + i]; |
| 474 | let x2 = x_data[x_off + half + i]; |
| 475 | out[x_off + i] = x1 * c - x2 * s; |
| 476 | out[x_off + half + i] = x2 * c + x1 * s; |
| 477 | } |
| 478 | } |
| 479 | return Tensor::from_vec(out, x_shape, x.device()); |
| 480 | } |
| 481 | candle_nn::rotary_emb::rope(x, cos, sin) |
| 482 | } |
| 483 | |
| 484 | /// SiLU (Swish) activation: `x * sigmoid(x)`. |
| 485 | fn silu(&self, x: &Tensor) -> Result<Tensor> { |