candding 0.1.0

Pure candle embeddings for Rust with no ONNX: dense, sparse, multi-vector and rerank models today, image models on the roadmap
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

//! RMSNorm in the two conventions the decoder families use.

use candle_core::Tensor;
use candle_nn::VarBuilder;

use crate::error::Result;

/// `Scale` multiplies by `w` (Llama, Qwen3); `OnePlusScale` multiplies by `1 + w` (Gemma3).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RmsNormStyle {
  Scale,
  OnePlusScale,
}

pub struct RmsNorm {
  weight: Tensor,
  eps: f32,
}

impl RmsNorm {
  /// `1 + w` is materialized at load so both styles share one forward path.
  /// The unit init hint only serves randomly initialised test models; real checkpoints carry the tensor.
  pub fn load(vb: VarBuilder, size: usize, eps: f64, style: RmsNormStyle) -> Result<Self> {
    let weight = vb.get_with_hints(size, "weight", candle_nn::Init::Const(1.0))?;
    let weight = match style {
      RmsNormStyle::Scale => weight,
      RmsNormStyle::OnePlusScale => (weight + 1.0)?,
    };
    Ok(Self {
      weight,
      eps: eps as f32,
    })
  }

  /// Normalizes the last dimension; the weight is cast per call because candle's `rms_norm` needs matching dtypes.
  pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
    let weight = self.weight.to_dtype(xs.dtype())?;
    Ok(candle_nn::ops::rms_norm(
      &xs.contiguous()?,
      &weight,
      self.eps,
    )?)
  }
}

#[cfg(test)]
mod tests {
  use candle_core::DType;
  use candle_core::Device;
  use candle_core::Tensor;
  use candle_nn::VarBuilder;
  use candle_nn::VarMap;

  use super::*;

  fn norm(style: RmsNormStyle, weight: &[f32]) -> RmsNorm {
    let mut varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
    RmsNorm::load(vb.pp("n"), weight.len(), 1e-6, RmsNormStyle::Scale).unwrap();
    varmap
      .set_one("n.weight", Tensor::new(weight, &Device::Cpu).unwrap())
      .unwrap();
    RmsNorm::load(vb.pp("n"), weight.len(), 1e-6, style).unwrap()
  }

  fn rows(t: Tensor) -> Vec<Vec<f32>> {
    t.to_vec2::<f32>().unwrap()
  }

  #[test]
  fn scale_style_multiplies_by_the_weight() {
    let xs = Tensor::new(&[[3f32, 4.0]], &Device::Cpu).unwrap();
    let out = rows(norm(RmsNormStyle::Scale, &[1.0, 2.0]).forward(&xs).unwrap());
    let rms = (12.5f32 + 1e-6).sqrt();
    assert!(
      (out[0][0] - 3.0 / rms).abs() < 1e-5 && (out[0][1] - 8.0 / rms).abs() < 1e-5,
      "{out:?}"
    );
  }

  #[test]
  fn one_plus_style_adds_one_to_the_weight() {
    let xs = Tensor::new(&[[3f32, 4.0]], &Device::Cpu).unwrap();
    let out = rows(
      norm(RmsNormStyle::OnePlusScale, &[0.0, 1.0])
        .forward(&xs)
        .unwrap(),
    );
    let rms = (12.5f32 + 1e-6).sqrt();
    assert!(
      (out[0][0] - 3.0 / rms).abs() < 1e-5 && (out[0][1] - 8.0 / rms).abs() < 1e-5,
      "{out:?}"
    );
  }

  #[test]
  fn half_precision_input_keeps_its_dtype() {
    let xs = Tensor::new(&[[3f32, 4.0]], &Device::Cpu)
      .unwrap()
      .to_dtype(DType::F16)
      .unwrap();
    let out = norm(RmsNormStyle::Scale, &[1.0, 1.0]).forward(&xs).unwrap();
    assert_eq!(out.dtype(), DType::F16);
  }
}