topos 0.13.1

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

use crate::graph::Network;
use crate::op::Op;
use crate::{Tape, Tensor, conv2d};

use super::super::pattern::Pattern;
use super::super::view::View;
use super::match_at;

/// Records a padded strided convolution and returns its network.
fn conv_network() -> Network<f64> {
    let tape = Tape::new();
    let input = tape.leaf(Tensor::new(
        [1, 2, 4, 4],
        (0..32).map(|v| v as f64 / 8.0 - 2.0).collect::<Vec<_>>(),
    ));
    let weights = tape.leaf(Tensor::new(
        [2, 2, 2, 2],
        (0..16).map(|v| v as f64 / 4.0 - 2.0).collect::<Vec<_>>(),
    ));
    let bias = tape.leaf(Tensor::new([2], [0.25, -0.5]));
    let _output = conv2d(input, weights, bias, 2, 1);
    tape.into_network()
}

/// Returns the index of the network's single `matmul` node.
fn matmul_index(network: &Network<f64>) -> usize {
    let ops = &network.structure().ops;
    (0..ops.len())
        .find(|&index| matches!(ops.get(index), Some(Op::MatMul(_))))
        .expect("the conv records one matmul")
}

#[test]
fn the_conv_chain_matches_with_its_geometry() {
    let network = conv_network();
    let structure = network.structure();
    let length = structure.ops.len();
    let wanted = vec![true; length];
    let readable = vec![false; length];
    let view = View::new(structure, &wanted, &readable);
    let matmul = matmul_index(&network);

    let candidate = match_at(matmul, &view).expect("the canonical chain matches");
    let group = match candidate.pattern {
        Pattern::WindowProduct(group) => group,
        _ => panic!("the conv chain matches as a window product"),
    };
    assert_eq!(structure.shapes[group.source].axes(), [1, 2, 4, 4]);
    assert_eq!(structure.shapes[group.kernel].axes(), [8, 2]);
    assert_eq!(group.kernel_height, 2);
    assert_eq!(group.kernel_width, 2);
    assert_eq!(group.stride, 2);
    assert_eq!(group.padding, 1);
    // The reshape, permute, two unfolds, and two symmetric pads.
    assert_eq!(candidate.interiors.len(), 6);
    assert!(candidate.named.is_empty());
}

#[test]
fn a_kept_interior_bars_the_match() {
    let network = conv_network();
    let structure = network.structure();
    let length = structure.ops.len();
    let matmul = matmul_index(&network);
    let patches = structure
        .operands
        .get(matmul)
        .expect("plan columns are fixed")
        .as_slice()[0]
        .index();

    let wanted = vec![true; length];
    let mut readable = vec![false; length];
    readable[patches] = true;
    let view = View::new(structure, &wanted, &readable);
    assert!(match_at(matmul, &view).is_none());
}