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

use crate::{Shape, Tape, Tensor};

use super::super::{Module, init};
use super::Dropout;

#[test]
fn masks_are_deterministic_in_seed_and_step() {
    let mut first = init::dropout::<f64>(11, 0.75);
    let mut second = init::dropout::<f64>(11, 0.75);
    let shape = Shape::new([4, 8]);
    for _ in 0..3 {
        assert_eq!(first(&shape).to_vec(), second(&shape).to_vec());
    }
    let mut third = init::dropout::<f64>(12, 0.75);
    let mut fresh = init::dropout::<f64>(11, 0.75);
    assert_ne!(fresh(&shape).to_vec(), third(&shape).to_vec());
}

#[test]
fn masks_hold_the_inverted_scale() {
    let mut masks = init::dropout::<f64>(7, 0.8);
    let mask = masks(&Shape::new([64, 64])).to_vec();
    assert!(mask.iter().all(|&value| value == 0.0 || value == 1.0 / 0.8));
    // Deterministic for the fixed seed, and near the keep probability
    // over 4096 draws.
    let kept = mask.iter().filter(|&&value| value != 0.0).count();
    assert!((2900..3600).contains(&kept), "kept {kept} of 4096");
}

#[test]
#[should_panic(expected = "the keep probability must lie within (0, 1], got 0")]
fn a_zero_keep_probability_is_rejected() {
    let _ = init::dropout::<f64>(1, 0.0);
}

#[test]
fn an_unfed_run_is_the_identity() {
    // The mask input's default payload is all ones, so the masked
    // expression answers bitwise as the plain one — inference needs
    // no second expression and no mode flag.
    let tape: Tape<f64> = Tape::new();
    let x = tape.parameter(Tensor::new(
        [2, 3],
        vec![0.5_f64, -1.0, 2.0, 3.0, -0.25, 1.5],
    ));
    let dropout = Dropout::new(&tape, [2, 3]);
    let masked = dropout.express(x + x).symbol();
    let plain = (x + x).symbol();
    let network = tape.into_network();
    let run = network.forward(&network.parameters(), []);
    assert_eq!(run.of(masked).to_vec(), run.of(plain).to_vec());
}

#[test]
fn gradients_route_through_the_mask() {
    let tape = Tape::new();
    let x = tape.parameter(Tensor::new([1, 4], [1.0_f64, 2.0, 3.0, 4.0]));
    let dropout = Dropout::new(&tape, [1, 4]);
    let loss = dropout.express(x).sum();
    let (loss, x) = (loss.symbol(), x.symbol());
    let network = tape.into_network();
    let mask = Tensor::new([1, 4], [2.0_f64, 0.0, 2.0, 0.0]);
    let run = network.forward(&network.parameters(), [(dropout.mask(), mask)]);
    assert_eq!(run.backward(loss).of(x).to_vec(), vec![2.0, 0.0, 2.0, 0.0]);
}

#[test]
fn training_replays_bitwise_under_matched_seeds() {
    // A tiny masked training loop, run twice from the same seeds:
    // the trajectories agree bit for bit, dropout masks included.
    fn final_loss() -> f64 {
        let tape = Tape::new();
        let w = tape.parameter(Tensor::new([2, 2], [0.5_f64, -0.25, 0.75, 0.1]));
        let x = tape.leaf(Tensor::new([2, 2], [1.0_f64, 2.0, -1.0, 0.5]));
        let dropout = Dropout::new(&tape, [2, 2]);
        let product = dropout.express(x.matmul(w));
        let loss = (product * product).sum().symbol();
        let network = tape.into_network();

        let mut masks = init::dropout::<f64>(21, 0.5);
        let rate = Tensor::new([], [0.05_f64]);
        let mut parameters = network.parameters();
        let mut last = 0.0;
        for _ in 0..8 {
            let run = network.forward(&parameters, [(dropout.mask(), masks(&Shape::new([2, 2])))]);
            last = run.of(loss).scalar();
            let gradients = run.backward(loss).parameters(&parameters);
            parameters = parameters.step(&gradients, |parameter, gradient| {
                parameter.clone() - gradient.clone() * rate.broadcast_like(gradient)
            });
        }
        last
    }
    assert_eq!(final_loss().to_bits(), final_loss().to_bits());
}

#[test]
fn dropout_composes_as_a_module() {
    // The trait impl serves the record-time composition surface: an
    // object-safe module slot expresses identically to the inherent
    // method.
    let tape = Tape::new();
    let x = tape.parameter(Tensor::new([2, 2], [1.0_f64, -2.0, 3.0, -4.0]));
    let dropout = Dropout::new(&tape, [2, 2]);
    let module: &dyn Module<f64> = &dropout;
    let through_trait = module.express(x).symbol();
    let network = tape.into_network();
    let mask = Tensor::new([2, 2], [2.0_f64, 0.0, 0.0, 2.0]);
    let run = network.forward(&network.parameters(), [(dropout.mask(), mask)]);
    assert_eq!(run.of(through_trait).to_vec(), vec![2.0, 0.0, 0.0, -8.0]);
}