topos 0.13.0

An autodiff compiler stack in Rust.
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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243

use crate::{Adam, AdamW, Entry, Optimizer, Sgd, Tape, Tensor};

/// The conventional hyperparameters as single-value payloads.
fn conventional() -> (Tensor<f64>, Tensor<f64>, Tensor<f64>) {
    (
        Tensor::new([], [0.9_f64]),
        Tensor::new([], [0.999_f64]),
        Tensor::new([], [1e-8_f64]),
    )
}

#[test]
fn two_steps_match_the_paper_trace() {
    // `loss = w * w` over a single scalar-shaped parameter: the
    // gradient is `2w`, and two Adam steps are traced by hand from
    // the paper's update rule.
    let tape = Tape::new();
    let w = tape.parameter(Tensor::new([], [1.0_f64]));
    let loss = (w * w).sum().symbol();
    let w = w.symbol();
    let network = tape.into_network();
    let (beta1, beta2, epsilon) = conventional();
    let mut adam = Adam::new(beta1, beta2, epsilon);
    let rate = Tensor::new([], [0.1_f64]);

    let mut parameters = network.parameters();
    let mut expected_w = 1.0_f64;
    let (mut m, mut v) = (0.0_f64, 0.0);
    let (mut beta1_power, mut beta2_power) = (1.0_f64, 1.0);
    for _ in 0..2 {
        let run = network.forward(&parameters, []);
        let gradients = run.backward(loss).parameters(&parameters);
        parameters = adam.step(&parameters, &gradients, &rate);

        let gradient = 2.0 * expected_w;
        m = m * 0.9 + gradient * (1.0 - 0.9);
        v = v * 0.999 + gradient * gradient * (1.0 - 0.999);
        beta1_power *= 0.9;
        beta2_power *= 0.999;
        let corrected_m = m / (1.0 - beta1_power);
        let corrected_v = v / (1.0 - beta2_power);
        expected_w -= 0.1 * corrected_m / (corrected_v.sqrt() + 1e-8);

        let stepped = parameters.of(w).to_vec()[0];
        assert!(
            (stepped - expected_w).abs() < 1e-12,
            "stepped {stepped}, expected {expected_w}"
        );
    }
}

#[test]
fn identical_runs_are_bitwise_identical() {
    let run = || {
        let tape = Tape::new();
        let w = tape.parameter(Tensor::new([2, 2], [1.0_f64, -0.5, 0.25, 2.0]));
        let loss = (w * w).sum().symbol();
        let w = w.symbol();
        let network = tape.into_network();
        let (beta1, beta2, epsilon) = conventional();
        let mut adam = Adam::new(beta1, beta2, epsilon);
        let rate = Tensor::new([], [0.05_f64]);
        let mut parameters = network.parameters();
        for _ in 0..5 {
            let run = network.forward(&parameters, []);
            let gradients = run.backward(loss).parameters(&parameters);
            parameters = adam.step(&parameters, &gradients, &rate);
        }
        parameters.of(w).to_vec()
    };
    for (first, second) in run().iter().zip(run()) {
        assert_eq!(first.to_bits(), second.to_bits());
    }
}

#[test]
fn adamw_decays_weights_and_spares_biases() {
    let tape = Tape::new();
    let weights = tape.parameter(Tensor::new([2, 2], [1.0_f64, -0.5, 0.25, 2.0]));
    let bias = tape.parameter(Tensor::new([2], [0.5_f64, -1.0]));
    let x = tape.leaf(Tensor::new([1, 2], [1.0_f64, -1.0]));
    let product = x.matmul(weights);
    let shifted = product + bias.broadcast_along_like(0, product);
    let loss = (shifted * shifted).sum();
    let (weights, bias, loss) = (weights.symbol(), bias.symbol(), loss.symbol());
    let network = tape.into_network();

    let (beta1, beta2, epsilon) = conventional();
    let mut plain = Adam::new(beta1.clone(), beta2.clone(), epsilon.clone());
    let mut decoupled = AdamW::new(beta1, beta2, epsilon, Tensor::new([], [0.1_f64]));
    let rate = Tensor::new([], [0.05_f64]);

    let parameters = network.parameters();
    let run = network.forward(&parameters, []);
    let gradients = run.backward(loss).parameters(&parameters);
    let by_adam = plain.step(&parameters, &gradients, &rate);
    let by_adamw = decoupled.step(&parameters, &gradients, &rate);

    // The rank-1 bias is spared: both routes agree bitwise there.
    let adam_bias = by_adam.of(bias).to_vec();
    let adamw_bias = by_adamw.of(bias).to_vec();
    for (plain, decayed) in adam_bias.iter().zip(&adamw_bias) {
        assert_eq!(plain.to_bits(), decayed.to_bits());
    }

    // The rank-2 weight differs by exactly the decoupled decay term.
    let before = parameters.of(weights).to_vec();
    let adam_weights = by_adam.of(weights).to_vec();
    let adamw_weights = by_adamw.of(weights).to_vec();
    for ((plain, decayed), original) in adam_weights.iter().zip(&adamw_weights).zip(before) {
        let term = original * 0.1 * 0.05;
        assert!((plain - decayed - term).abs() < 1e-15);
    }
}

#[test]
fn step_where_overrides_the_structural_policy() {
    let tape = Tape::new();
    let weights = tape.parameter(Tensor::new([2, 2], [1.0_f64, -0.5, 0.25, 2.0]));
    let loss = (weights * weights).sum().symbol();
    let weights = weights.symbol();
    let network = tape.into_network();
    let (beta1, beta2, epsilon) = conventional();
    let mut plain = Adam::new(beta1.clone(), beta2.clone(), epsilon.clone());
    let mut decoupled = AdamW::new(beta1, beta2, epsilon, Tensor::new([], [0.1_f64]));
    let rate = Tensor::new([], [0.05_f64]);

    let parameters = network.parameters();
    let run = network.forward(&parameters, []);
    let gradients = run.backward(loss).parameters(&parameters);
    let by_adam = plain.step(&parameters, &gradients, &rate);
    // Decay nothing: AdamW must reproduce Adam bitwise.
    let spared = decoupled.step_where(&parameters, &gradients, &rate, |_, _| false);

    for (plain, decayed) in by_adam
        .of(weights)
        .to_vec()
        .iter()
        .zip(spared.of(weights).to_vec())
    {
        assert_eq!(plain.to_bits(), decayed.to_bits());
    }
}

#[test]
fn recorded_gradients_feed_adam_bitwise() {
    // Two grains, one trajectory: recorded gradients arrive
    // parameter-aligned, an engine backward projects onto the same
    // slots, and Adam cannot tell the routes apart.
    let build = || {
        let tape = Tape::new();
        let w = tape.parameter(Tensor::new([2, 2], [1.0_f64, -0.5, 0.25, 2.0]));
        let loss = (w * w).sum();
        let (w, loss) = (w.symbol(), loss.symbol());
        (tape, w, loss)
    };
    let (beta1, beta2, epsilon) = conventional();
    let rate = Tensor::new([], [0.05_f64]);

    let (engine_tape, engine_w, engine_loss) = build();
    let engine_network = engine_tape.into_network();
    let mut engine_adam = Adam::new(beta1.clone(), beta2.clone(), epsilon.clone());
    let mut engine_parameters = engine_network.parameters();

    let (recorded_tape, recorded_w, recorded_loss) = build();
    let adjoints = recorded_tape.differentiate(recorded_loss, [recorded_w]);
    let recorded_network = recorded_tape.into_network();
    let plan = recorded_network.compile(Entry::roots(adjoints.roots()));
    let mut recorded_adam = Adam::new(beta1, beta2, epsilon);
    let mut recorded_parameters = recorded_network.parameters();

    for _ in 0..4 {
        let run = engine_network.forward(&engine_parameters, []);
        let gradients = run.backward(engine_loss).parameters(&engine_parameters);
        engine_parameters = engine_adam.step(&engine_parameters, &gradients, &rate);

        let run = plan.forward(&recorded_parameters, []);
        let gradients = run.recorded_gradients(&adjoints);
        recorded_parameters = recorded_adam.step(&recorded_parameters, &gradients, &rate);

        for (engine, recorded) in engine_parameters
            .of(engine_w)
            .to_vec()
            .iter()
            .zip(recorded_parameters.of(recorded_w).to_vec())
        {
            assert_eq!(engine.to_bits(), recorded.to_bits());
        }
    }
}

#[test]
fn moments_cover_parameter_slots_not_graph_nodes() {
    let tape = Tape::new();
    let w = tape.parameter(Tensor::new([2, 2], [1.0_f64, -0.5, 0.25, 2.0]));
    let loss = (w * w).sum().symbol();
    let network = tape.into_network();
    let (beta1, beta2, epsilon) = conventional();
    let mut adam = Adam::new(beta1, beta2, epsilon);
    let rate = Tensor::new([], [0.05_f64]);
    let parameters = network.parameters();

    let run = network.forward(&parameters, []);
    let gradients = run.backward(loss).parameters(&parameters);
    adam.step(&parameters, &gradients, &rate);

    // The moment tables are parameter-aligned: one entry per slot,
    // never one per recorded node.
    let first = adam.first.as_ref().expect("a step ran");
    assert_eq!(first.len(), parameters.len());
    assert!(network.len() > parameters.len());
}

#[test]
fn adam_descends_faster_than_sgd_on_a_skewed_bowl() {
    // A quadratic bowl with wildly different curvatures per axis: the
    // fixed problem where per-coordinate step normalization pays.
    let run = |strategy: &mut dyn Optimizer<f64>| {
        let tape = Tape::new();
        let w = tape.parameter(Tensor::new([2], [5.0_f64, 5.0]));
        let curvatures = tape.leaf(Tensor::new([2], [100.0_f64, 0.01]));
        let loss = (w * w * curvatures).sum().symbol();
        let network = tape.into_network();
        let rate = Tensor::new([], [0.01_f64]);
        let mut parameters = network.parameters();
        for _ in 0..100 {
            let run = network.forward(&parameters, []);
            let gradients = run.backward(loss).parameters(&parameters);
            parameters = strategy.step(&parameters, &gradients, &rate);
        }
        let run = network.forward(&parameters, []);
        run.of(loss).scalar()
    };

    let (beta1, beta2, epsilon) = conventional();
    let sgd_loss = run(&mut Sgd);
    let adam_loss = run(&mut Adam::new(beta1, beta2, epsilon));
    assert!(adam_loss.is_finite() && sgd_loss.is_finite());
    assert!(
        adam_loss < sgd_loss,
        "adam {adam_loss} should beat sgd {sgd_loss} on the skewed bowl"
    );
}