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

use super::*;

/// On CPU both functions run the same `matmul_attention`.
#[test]
fn scaled_dot_product_and_the_weighted_sibling_agree_on_cpu() {
  let q = Tensor::randn(0f32, 1.0, (2, 4, 3, 8), &Device::Cpu).unwrap();
  let k = Tensor::randn(0f32, 1.0, (2, 4, 3, 8), &Device::Cpu).unwrap();
  let v = Tensor::randn(0f32, 1.0, (2, 4, 3, 8), &Device::Cpu).unwrap();
  let attention_mask = Tensor::new(&[[1u32, 1, 0], [1, 1, 1]], &Device::Cpu).unwrap();
  let mask = AttentionMask::padding(&attention_mask, &[2, 3], 4, 8, DType::F32).unwrap();
  let plain = scaled_dot_product(&q, &k, &v, &mask, 0.5).unwrap();
  let (with_weights, _weights) = scaled_dot_product_with_weights(&q, &k, &v, &mask, 0.5).unwrap();
  let diff = max_diff(&plain, &with_weights);
  assert!(diff < 1e-6, "cpu context max abs diff {diff}");
}

/// On Metal the context comes from the fused kernel and the weights from the matmul path, so they agree only loosely;
/// rows summing to one show the sibling stayed off the fused path, which has no softmax to hand back.
#[cfg(feature = "metal")]
#[test]
fn scaled_dot_product_and_the_weighted_sibling_agree_on_metal() {
  let device = Device::new_metal(0).expect("a metal device");
  let head_dim = 64;
  assert!(metal_sdpa_applies(&device, head_dim, DType::F32));
  let q = Tensor::randn(0f32, 1.0, (2, 4, 5, head_dim), &device).unwrap();
  let k = Tensor::randn(0f32, 1.0, (2, 4, 5, head_dim), &device).unwrap();
  let v = Tensor::randn(0f32, 1.0, (2, 4, 5, head_dim), &device).unwrap();
  let scale = 1.0 / (head_dim as f64).sqrt();
  let fused = scaled_dot_product(&q, &k, &v, &AttentionMask::none(), scale).unwrap();
  let (matmul, weights) =
    scaled_dot_product_with_weights(&q, &k, &v, &AttentionMask::none(), scale).unwrap();
  let diff = max_diff(
    &fused.to_device(&Device::Cpu).unwrap(),
    &matmul.to_device(&Device::Cpu).unwrap(),
  );
  assert!(diff < 1e-4, "fused vs matmul max abs diff {diff}");
  let sums = values(&weights.sum(3).unwrap().to_device(&Device::Cpu).unwrap());
  for sum in sums {
    assert!((sum - 1.0).abs() < 1e-4, "{sum}");
  }
}

/// `v` is not the identity, which would make the context equal the weights and hide a swap of the two.
#[test]
fn weights_match_hand_computation() {
  let q = t(&[1.0, 0.0, 0.0, 1.0], (1, 1, 2, 2));
  let k = t(&[0.0, 0.0, 3f32.ln(), 0.0], (1, 1, 2, 2));
  let v = t(&[2.0, 3.0, 5.0, 7.0], (1, 1, 2, 2));
  let (_, weights) =
    scaled_dot_product_with_weights(&q, &k, &v, &AttentionMask::none(), 1.0).unwrap();
  let rows = weights
    .squeeze(0)
    .unwrap()
    .squeeze(0)
    .unwrap()
    .to_vec2::<f32>()
    .unwrap();
  assert!(
    (rows[0][0] - 0.25).abs() < 1e-6 && (rows[0][1] - 0.75).abs() < 1e-6,
    "{rows:?}"
  );
  assert!(
    (rows[1][0] - 0.5).abs() < 1e-6 && (rows[1][1] - 0.5).abs() < 1e-6,
    "{rows:?}"
  );
}

#[test]
fn each_row_of_the_weights_sums_to_one() {
  let q = Tensor::randn(0f32, 1.0, (2, 3, 4, 5), &Device::Cpu).unwrap();
  let k = Tensor::randn(0f32, 1.0, (2, 3, 4, 5), &Device::Cpu).unwrap();
  let v = Tensor::randn(0f32, 1.0, (2, 3, 4, 5), &Device::Cpu).unwrap();
  let (_, weights) =
    scaled_dot_product_with_weights(&q, &k, &v, &AttentionMask::none(), 0.4).unwrap();
  let sums = values(&weights.sum(3).unwrap());
  for sum in sums {
    assert!((sum - 1.0).abs() < 1e-5, "{sum}");
  }
}

#[test]
fn a_padded_position_contributes_nothing_to_a_real_positions_row() {
  let q = t(&[1.0, 0.0], (1, 1, 1, 2));
  let k = t(&[0.0, 0.0, 3f32.ln(), 0.0], (1, 1, 2, 2));
  let v = t(&[9.0, 4.0, 6.0, 2.0], (1, 1, 2, 2));
  let attention_mask = Tensor::new(&[[1u32, 0]], &Device::Cpu).unwrap();
  let mask = AttentionMask::padding(&attention_mask, &[1], 1, 2, DType::F32).unwrap();
  let (_, weights) = scaled_dot_product_with_weights(&q, &k, &v, &mask, 1.0).unwrap();
  let row = values(&weights);
  assert!(
    (row[0] - 1.0).abs() < 1e-6 && row[1].abs() < 1e-6,
    "{row:?}"
  );
}