candding 0.1.0

//! Rotary position embeddings in the rotate-half convention of Llama, Qwen3 and Gemma3.

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;
use candle_nn::rotary_emb::rope;

use crate::error::Result;

pub struct RotaryEmbedding {
  cos: Tensor,
  sin: Tensor,
}

impl RotaryEmbedding {
  /// Tables for `0..max_positions`, computed in F32 and stored in `dtype`; `inv_freq[i] = theta^(-2i/head_dim)`.
  pub fn new(
    head_dim: usize,
    max_positions: usize,
    theta: f64,
    dtype: DType,
    device: &Device,
  ) -> Result<Self> {
    Self::with_frequency_scale(head_dim, max_positions, theta, 1.0, dtype, device)
  }

  /// Like `new`, with every `inv_freq[i]` also multiplied by `freq_scale` (gte_new's NTK scaling).
  pub fn with_frequency_scale(
    head_dim: usize,
    max_positions: usize,
    theta: f64,
    freq_scale: f64,
    dtype: DType,
    device: &Device,
  ) -> Result<Self> {
    let half = head_dim / 2;
    let inv_freq: Vec<f32> = (0..half)
      .map(|i| (freq_scale / theta.powf(2.0 * i as f64 / head_dim as f64)) as f32)
      .collect();
    let inv_freq = Tensor::from_vec(inv_freq, (1, half), device)?;
    let positions = Tensor::arange(0u32, max_positions as u32, device)?
      .to_dtype(DType::F32)?
      .reshape((max_positions, 1))?;
    let freqs = positions.matmul(&inv_freq)?;
    Ok(Self {
      cos: freqs.cos()?.to_dtype(dtype)?,
      sin: freqs.sin()?.to_dtype(dtype)?,
    })
  }

  /// `q` and `k` are `(batch, heads, seq, head_dim)`; positions run `0..seq` in every row, true under right padding.
  pub fn apply(&self, q: &Tensor, k: &Tensor) -> Result<(Tensor, Tensor)> {
    let seq = q.dim(2)?;
    let cos = self.cos.narrow(0, 0, seq)?;
    let sin = self.sin.narrow(0, 0, seq)?;
    Ok((
      rope(&q.contiguous()?, &cos, &sin)?,
      rope(&k.contiguous()?, &cos, &sin)?,
    ))
  }

  pub fn table(&self) -> (&Tensor, &Tensor) {
    (&self.cos, &self.sin)
  }
}

#[cfg(test)]
mod tests {
  use std::fs;
  use std::path::PathBuf;

  use candle_core::Device;
  use serde::Deserialize;

  use super::*;

  #[derive(Deserialize)]
  struct Fixture {
    head_dim: usize,
    positions: usize,
    cases: Vec<Case>,
  }

  #[derive(Deserialize)]
  struct Case {
    theta: f64,
    cos: Vec<Vec<f32>>,
    sin: Vec<Vec<f32>>,
  }

  fn fixture() -> Fixture {
    let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/rope.json");
    serde_json::from_str(&fs::read_to_string(path).unwrap()).unwrap()
  }

  fn max_diff(a: &[Vec<f32>], b: &[Vec<f32>]) -> f32 {
    a.iter()
      .zip(b)
      .flat_map(|(x, y)| x.iter().zip(y).map(|(p, q)| (p - q).abs()))
      .fold(0.0, f32::max)
  }

  #[test]
  fn tables_match_the_python_fixture() {
    let fx = fixture();
    for case in &fx.cases {
      let rope = RotaryEmbedding::new(
        fx.head_dim,
        fx.positions,
        case.theta,
        DType::F32,
        &Device::Cpu,
      )
      .unwrap();
      let (cos, sin) = rope.table();
      assert!(
        max_diff(&cos.to_vec2::<f32>().unwrap(), &case.cos) < 1e-6,
        "cos theta {}",
        case.theta
      );
      assert!(
        max_diff(&sin.to_vec2::<f32>().unwrap(), &case.sin) < 1e-6,
        "sin theta {}",
        case.theta
      );
    }
  }

  #[test]
  fn apply_matches_the_slow_reference_and_leaves_position_zero_alone() {
    let rope = RotaryEmbedding::new(8, 16, 10000.0, DType::F32, &Device::Cpu).unwrap();
    let q = Tensor::randn(0f32, 1.0, (2, 3, 5, 8), &Device::Cpu).unwrap();
    let k = Tensor::randn(0f32, 1.0, (2, 1, 5, 8), &Device::Cpu).unwrap();
    let (rq, rk) = rope.apply(&q, &k).unwrap();
    let (cos, sin) = rope.table();
    let cos = cos.narrow(0, 0, 5).unwrap();
    let sin = sin.narrow(0, 0, 5).unwrap();
    let slow_q = candle_nn::rotary_emb::rope_slow(&q, &cos, &sin).unwrap();
    let slow_k = candle_nn::rotary_emb::rope_slow(&k, &cos, &sin).unwrap();
    let diff = |a: &Tensor, b: &Tensor| {
      (a - b)
        .unwrap()
        .abs()
        .unwrap()
        .flatten_all()
        .unwrap()
        .max(0)
        .unwrap()
        .to_scalar::<f32>()
        .unwrap()
    };
    assert!(diff(&rq, &slow_q) < 1e-5 && diff(&rk, &slow_k) < 1e-5);
    let first = |t: &Tensor| t.narrow(2, 0, 1).unwrap();
    assert!(
      diff(&first(&rq), &first(&q)) < 1e-6,
      "position 0 must be unchanged"
    );
    assert_eq!(rq.dims(), q.dims());
    assert_eq!(rk.dims(), k.dims());
  }

  /// The reference's loaded `inv_freq`: 0.9370838 and 0.6443913 for gte-multilingual-base, 0.9785720 and 0.6354667 for gte-base-en-v1.5.
  #[test]
  fn a_frequency_scale_reproduces_the_reference_tables() {
    for (theta, factor, expected) in [
      (20_000.0f64, 8.0f64, [0.9370838f32, 0.6443913]),
      (500_000.0, 2.0, [0.978_572, 0.6354667]),
    ] {
      let scale = factor.powf(-2.0 / 64.0);
      let rope = RotaryEmbedding::with_frequency_scale(
        64,
        4,
        theta * factor,
        scale,
        DType::F32,
        &Device::Cpu,
      )
      .unwrap();
      // cos(1 * inv_freq[i]) at position 1 recovers inv_freq[i] through acos for these magnitudes.
      let cos = rope.table().0.to_vec2::<f32>().unwrap();
      for (i, want) in expected.iter().enumerate() {
        let got = cos[1][i].acos();
        assert!(
          (got - want).abs() < 1e-5,
          "theta {theta} factor {factor} entry {i}: {got} vs {want}"
        );
      }
    }
  }

  #[test]
  fn new_is_the_unscaled_case() {
    let plain = RotaryEmbedding::new(64, 4, 10_000.0, DType::F32, &Device::Cpu).unwrap();
    let scaled =
      RotaryEmbedding::with_frequency_scale(64, 4, 10_000.0, 1.0, DType::F32, &Device::Cpu)
        .unwrap();
    assert_eq!(
      plain.table().0.to_vec2::<f32>().unwrap(),
      scaled.table().0.to_vec2::<f32>().unwrap()
    );
  }
}