legume-numeric 0.8.11

Numeric and ML foundation for the legume ecosystem (matrix, Leiden, candle, MCMC)
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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

use super::RowAdagrad;
use candle_core::{DType, Device, Tensor, Var};

fn approx(a: f32, b: f32, tol: f32) -> bool {
    (a - b).abs() <= tol
}

#[test]
fn one_row_adagrad_step_matches_the_closed_form() {
    let dev = Device::Cpu;
    let var =
        Var::from_tensor(&Tensor::from_vec(vec![1f32, 2., 3., 4., 5., 6.], (3, 2), &dev).unwrap())
            .unwrap();
    let grad = Tensor::from_vec(vec![1f32, -1., 0., 0., 2., 2.], (3, 2), &dev).unwrap();
    let mut opt = RowAdagrad::new(3, 0.1, &dev).unwrap();
    opt.step(&var, &grad).unwrap();
    // acc = mean_d(grad²) = [1, 0, 4]; std = sqrt(acc) + 1e-10
    // row0: p − 0.1·g/1 = [0.9, 2.1]; row2: p − 0.1·2/2 = [4.9, 5.9]
    let p = var.as_tensor().to_vec2::<f32>().unwrap();
    assert!(approx(p[0][0], 0.9, 1e-6) && approx(p[0][1], 2.1, 1e-6));
    assert!(approx(p[2][0], 4.9, 1e-6) && approx(p[2][1], 5.9, 1e-6));
    let acc = opt.accumulator().to_vec1::<f32>().unwrap();
    assert!(approx(acc[0], 1.0, 1e-6) && approx(acc[2], 4.0, 1e-6));
    // A second identical step divides by sqrt(2·acc).
    opt.step(&var, &grad).unwrap();
    let p = var.as_tensor().to_vec2::<f32>().unwrap();
    assert!(approx(p[0][0], 0.9 - 0.1 / 2f32.sqrt(), 1e-6));
    assert!(approx(p[2][1], 5.9 - 0.1 * 2.0 / 8f32.sqrt(), 1e-6));
}

#[test]
fn rows_with_zero_gradient_are_left_unchanged_and_keep_a_zero_accumulator() {
    let dev = Device::Cpu;
    let var = Var::from_tensor(
        &Tensor::from_vec(vec![1f32, 2., 3.3333, 4.25, 5., 6.], (3, 2), &dev).unwrap(),
    )
    .unwrap();
    let before = var.as_tensor().to_vec2::<f32>().unwrap();
    let grad = Tensor::from_vec(vec![1f32, -1., 0., 0., 2., 2.], (3, 2), &dev).unwrap();
    let mut opt = RowAdagrad::new(3, 0.1, &dev).unwrap();
    opt.step(&var, &grad).unwrap();
    opt.step(&var, &grad).unwrap();
    let after = var.as_tensor().to_vec2::<f32>().unwrap();
    assert_eq!(after[1], before[1], "bit-identical untouched row");
    assert_eq!(opt.accumulator().to_vec1::<f32>().unwrap()[1], 0.0);
    assert_ne!(after[0], before[0]);
}

#[test]
fn the_dense_gradient_of_index_select_sums_duplicate_rows_like_a_coalesced_sparse_update() {
    let dev = Device::Cpu;
    let var = Var::from_tensor(&Tensor::ones((6, 2), DType::F32, &dev).unwrap()).unwrap();
    let idx = Tensor::from_vec(vec![2u32, 2, 5], 3, &dev).unwrap();
    let w = Tensor::from_vec(vec![1f32, 1., 2., 2., 3., 3.], (3, 2), &dev).unwrap();
    let loss = var
        .as_tensor()
        .index_select(&idx, 0)
        .unwrap()
        .mul(&w)
        .unwrap()
        .sum_all()
        .unwrap();
    let grads = loss.backward().unwrap();
    let g = grads
        .get(&var)
        .expect("dense gradient on the table")
        .to_vec2::<f32>()
        .unwrap();
    assert_eq!(g.len(), 6);
    assert_eq!(g[2], vec![3.0, 3.0], "row 2 appears twice: 1 + 2");
    assert_eq!(g[5], vec![3.0, 3.0]);
    for r in [0, 1, 3, 4] {
        assert_eq!(g[r], vec![0.0, 0.0], "untouched row {r} is exactly zero");
    }
}

/// The accumulator is updated from the gradient tensor. If that update kept
/// its autograd history, the accumulator would hold a reference to every
/// step's gradient storage for the whole run (measured before the fix: RSS
/// grew ~650 MB per epoch on a 3000 × 2000 toy). Observable here with a
/// gradient whose history reaches a variable: after the step, nothing
/// reachable from the accumulator may lead back to it.
#[test]
fn the_accumulator_does_not_retain_the_gradients_autograd_graph() {
    let dev = Device::Cpu;
    let var =
        Var::from_tensor(&Tensor::from_vec(vec![1f32, 2., 3., 4., 5., 6.], (3, 2), &dev).unwrap())
            .unwrap();
    let source = Var::from_tensor(
        &Tensor::from_vec(vec![0.5f32, -1., 0., 0., 2., 1.], (3, 2), &dev).unwrap(),
    )
    .unwrap();
    // A "gradient" with a live history back to `source`.
    let g = source.as_tensor().affine(2.0, 0.0).unwrap();
    assert!(
        g.sum_all()
            .unwrap()
            .backward()
            .unwrap()
            .get(&source)
            .is_some(),
        "premise: the fake gradient is graph-attached"
    );
    let mut opt = RowAdagrad::new(3, 0.1, &dev).unwrap();
    opt.step(&var, &g).unwrap();
    let reach = opt.accumulator().sum_all().unwrap().backward().unwrap();
    assert!(
        reach.get(&source).is_none(),
        "the accumulator retains the gradient's graph"
    );
}

/// A row and its bias share one accumulator; a masked row keeps its value
/// while its bias moves with the bias-only accumulator; decay reaches only
/// the rows that carried a gradient.
#[test]
fn a_row_and_its_bias_share_the_accumulator_and_a_mask_pins_the_row_alone() {
    let dev = Device::Cpu;
    let row =
        Var::from_tensor(&Tensor::from_vec(vec![1f32, 2., 3., 4., 5., 6.], (3, 2), &dev).unwrap())
            .unwrap();
    let bias = Var::from_tensor(&Tensor::from_vec(vec![0f32, 0., 0.], 3, &dev).unwrap()).unwrap();
    let g_row = Tensor::from_vec(vec![1f32, -1., 0., 0., 2., 2.], (3, 2), &dev).unwrap();
    let g_bias = Tensor::from_vec(vec![1f32, 0., 3.], 3, &dev).unwrap();
    let mask = Tensor::from_vec(vec![1f32, 1., 0.], (3, 1), &dev).unwrap(); // row 2 pinned
    let mut opt = RowAdagrad::new(3, 0.1, &dev).unwrap();
    opt.step_with_bias(&row, &bias, &g_row, &g_bias, Some(&mask), 0.5)
        .unwrap();
    let p = row.as_tensor().to_vec2::<f32>().unwrap();
    let b = bias.as_tensor().to_vec1::<f32>().unwrap();
    // row 0: acc = (1 + 1 + 1) / 3 = 1 → step 0.1; decayed to half first
    assert!(approx(p[0][0], 0.5 - 0.1, 1e-6) && approx(p[0][1], 1.0 + 0.1, 1e-6));
    assert!(approx(b[0], -0.1, 1e-6));
    // row 1: no gradient at all — untouched, no decay, bias still
    assert_eq!(p[1], vec![3., 4.]);
    assert_eq!(b[1], 0.0);
    // row 2: pinned — the row keeps its value; the bias alone drives the
    // accumulator (9 / 1) so it moves by 0.1 · 3 / 3
    assert_eq!(p[2], vec![5., 6.]);
    assert!(approx(b[2], -0.1, 1e-6));
    let acc = opt.accumulator().to_vec1::<f32>().unwrap();
    assert!(approx(acc[0], 1.0, 1e-6) && acc[1] == 0.0 && approx(acc[2], 9.0, 1e-6));
}

/// A bias-only step: the accumulator takes the squared gradient itself (one
/// column, no mean), and the entry moves by `lr · g / sqrt(acc)` — what
/// `step_with_bias` does for the bias of a pinned row.
#[test]
fn a_bias_only_step_matches_the_closed_form() {
    let dev = Device::Cpu;
    let bias = Var::from_tensor(&Tensor::from_vec(vec![1f32, 2., 3.], 3, &dev).unwrap()).unwrap();
    let grad = Tensor::from_vec(vec![2f32, 0., -1.], 3, &dev).unwrap();
    let mut opt = RowAdagrad::new(3, 0.1, &dev).unwrap();
    opt.step_bias(&bias, &grad).unwrap();
    let b = bias.as_tensor().to_vec1::<f32>().unwrap();
    // acc = g² = [4, 0, 1]: entry 0 is 1 − 0.1·2/2, entry 2 is 3 + 0.1·1/1.
    assert!(approx(b[0], 0.9, 1e-6), "{b:?}");
    assert!(approx(b[1], 2.0, 1e-6), "{b:?}");
    assert!(approx(b[2], 3.1, 1e-6), "{b:?}");
    let acc = opt.accumulator().to_vec1::<f32>().unwrap();
    assert!(approx(acc[0], 4.0, 1e-6) && approx(acc[1], 0.0, 1e-6) && approx(acc[2], 1.0, 1e-6));
}