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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! Layer normalization that keeps its statistics in F32.
//! candle's fused kernel rounds the mean and variance to the input dtype, too coarse in BF16.

use candle_core::DType;
use candle_core::Module;
use candle_core::Result;
use candle_core::Tensor;
use candle_nn::VarBuilder;

pub struct LayerNorm {
  inner: candle_nn::LayerNorm,
}

impl LayerNorm {
  pub fn load(vb: VarBuilder, size: usize, eps: f64) -> Result<Self> {
    let config = candle_nn::LayerNormConfig {
      eps,
      remove_mean: true,
      affine: true,
    };
    let inner = candle_nn::layer_norm(size, config, vb.set_dtype(DType::F32))?;
    Ok(Self { inner })
  }

  /// For checkpoints with `norm_bias: false`, which carry no bias tensor.
  pub fn load_without_bias(vb: VarBuilder, size: usize, eps: f64) -> Result<Self> {
    let config = candle_nn::LayerNormConfig {
      eps,
      remove_mean: true,
      affine: false,
    };
    let inner = candle_nn::layer_norm(size, config, vb.set_dtype(DType::F32))?;
    Ok(Self { inner })
  }

  /// The output carries the input's dtype; only the normalization runs in F32.
  pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
    let dtype = xs.dtype();
    if dtype == DType::F32 {
      return self.inner.forward(xs);
    }
    self
      .inner
      .forward(&xs.to_dtype(DType::F32)?)?
      .to_dtype(dtype)
  }
}

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

  use super::*;

  fn unit_norm(dtype: DType) -> LayerNorm {
    let mut varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, dtype, &Device::Cpu);
    let norm = LayerNorm::load(vb, 4, 1e-12).unwrap();
    varmap
      .set_one("weight", Tensor::ones(4, DType::F32, &Device::Cpu).unwrap())
      .unwrap();
    varmap
      .set_one("bias", Tensor::zeros(4, DType::F32, &Device::Cpu).unwrap())
      .unwrap();
    norm
  }

  #[test]
  fn output_keeps_the_input_dtype_and_is_centred() {
    let norm = unit_norm(DType::F32);
    let xs = Tensor::new(&[[1.0f32, 2.0, 3.0, 4.0]], &Device::Cpu).unwrap();
    let out = norm.forward(&xs).unwrap();
    assert_eq!(out.dtype(), DType::F32);
    let row = out.to_vec2::<f32>().unwrap().remove(0);
    let mean: f32 = row.iter().sum::<f32>() / 4.0;
    assert!(mean.abs() < 1e-5, "{row:?}");
  }

  /// The F32-statistics path must beat a BF16 subtraction of two rounded values, the fused kernel's deviation.
  #[test]
  fn half_precision_tracks_the_f32_result() {
    let xs = Tensor::new(&[[10.0f32, 10.5, 11.0, 10.25]], &Device::Cpu).unwrap();
    let exact = unit_norm(DType::F32)
      .forward(&xs)
      .unwrap()
      .to_vec2::<f32>()
      .unwrap()
      .remove(0);
    let half = unit_norm(DType::F16)
      .forward(&xs.to_dtype(DType::F16).unwrap())
      .unwrap()
      .to_dtype(DType::F32)
      .unwrap()
      .to_vec2::<f32>()
      .unwrap()
      .remove(0);
    for (a, b) in exact.iter().zip(&half) {
      assert!((a - b).abs() < 2e-2, "{exact:?} vs {half:?}");
    }
  }

  /// ModernBERT's `norm_bias: false` norms carry no bias tensor, which `load` would ask for.
  #[test]
  fn load_without_bias_reads_the_weight_alone() {
    let mut tensors = std::collections::HashMap::new();
    tensors.insert(
      "weight".to_string(),
      Tensor::new(&[1f32, 1.0, 1.0, 1.0], &Device::Cpu).unwrap(),
    );
    let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
    assert!(LayerNorm::load(vb.clone(), 4, 1e-5).is_err());
    let norm = LayerNorm::load_without_bias(vb, 4, 1e-5).unwrap();
    let out = norm
      .forward(&Tensor::new(&[[1f32, 2.0, 3.0, 6.0]], &Device::Cpu).unwrap())
      .unwrap();
    let expected = [-1.069045, -0.534522, 0.0, 1.603567];
    for (a, e) in out.to_vec2::<f32>().unwrap()[0].iter().zip(expected) {
      assert!((a - e).abs() < 1e-4, "{a} vs {e}");
    }
  }
}