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

//! Gated MLP, `down(act(gate(x)) * up(x))`: SwiGLU for Qwen3, GeGLU for Gemma3, JinaBERT and gte_new.

use candle_core::Module;
use candle_core::Tensor;
use candle_nn::Linear;
use candle_nn::VarBuilder;
use candle_nn::linear_no_bias;

use crate::error::Result;
use crate::layers::activation::Activation;

pub struct GatedMlp {
  gate: Linear,
  up: Linear,
  down: Linear,
  activation: Activation,
}

impl GatedMlp {
  pub fn load(
    vb: VarBuilder,
    hidden: usize,
    intermediate: usize,
    activation: Activation,
  ) -> Result<Self> {
    Ok(Self::new(
      linear_no_bias(hidden, intermediate, vb.pp("gate_proj"))?,
      linear_no_bias(hidden, intermediate, vb.pp("up_proj"))?,
      linear_no_bias(intermediate, hidden, vb.pp("down_proj"))?,
      activation,
    ))
  }

  pub fn new(gate: Linear, up: Linear, down: Linear, activation: Activation) -> Self {
    Self {
      gate,
      up,
      down,
      activation,
    }
  }

  pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
    let gate = self.activation.forward(&self.gate.forward(xs)?)?;
    Ok(self.down.forward(&(gate * self.up.forward(xs)?)?)?)
  }
}

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

  use super::*;

  #[test]
  fn swiglu_with_identity_weights() {
    let eye = Tensor::new(&[[1f32, 0.0], [0.0, 1.0]], &Device::Cpu).unwrap();
    let mlp = GatedMlp::new(
      Linear::new(eye.clone(), None),
      Linear::new(eye.clone(), None),
      Linear::new(eye, None),
      Activation::Silu,
    );
    let out = mlp
      .forward(&Tensor::new(&[[1f32, 2.0]], &Device::Cpu).unwrap())
      .unwrap()
      .to_vec2::<f32>()
      .unwrap();
    let silu = |x: f32| x / (1.0 + (-x).exp());
    assert!(
      (out[0][0] - silu(1.0) * 1.0).abs() < 1e-5 && (out[0][1] - silu(2.0) * 2.0).abs() < 1e-5,
      "{out:?}"
    );
  }
}