tidu 0.2.0

Automatic-differentiation transforms (linearize, transpose, eager reverse-mode) for primitive computation graphs.
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

#[allow(unused_imports)]
use computegraph::types::{LocalValueId, OperationRole, ValueKey};
#[allow(unused_imports)]
use tidu::PrimitiveValue;

#[macro_export]
macro_rules! linearize_add {
    ($builder:expr, $OpAdd:path, $t0:expr, $t1:expr) => {
        Ok(match ($t0, $t1) {
            (Some(dx), Some(dy)) => {
                let out = $builder.add_primitive(
                    $OpAdd,
                    vec![PrimitiveValue::Local(dx), PrimitiveValue::Local(dy)],
                    OperationRole::Linearized {
                        active_mask: vec![true, true],
                    },
                );
                vec![Some(out[0])]
            }
            (Some(dx), None) => vec![Some(dx)],
            (None, Some(dy)) => vec![Some(dy)],
            (None, None) => vec![None],
        })
    };
}

#[macro_export]
macro_rules! linearize_mul {
    ($builder:expr, $OpMul:path, $OpAdd:path, $primal_in:expr, $t0:expr, $t1:expr) => {{
        let mut terms = Vec::new();
        if let Some(dx) = $t0 {
            let term = $builder.add_primitive(
                $OpMul,
                vec![
                    PrimitiveValue::Local(dx),
                    PrimitiveValue::External($primal_in[1].clone()),
                ],
                OperationRole::Linearized {
                    active_mask: vec![true, false],
                },
            );
            terms.push(term[0]);
        }
        if let Some(dy) = $t1 {
            let term = $builder.add_primitive(
                $OpMul,
                vec![
                    PrimitiveValue::External($primal_in[0].clone()),
                    PrimitiveValue::Local(dy),
                ],
                OperationRole::Linearized {
                    active_mask: vec![false, true],
                },
            );
            terms.push(term[0]);
        }
        Ok(match terms.as_slice() {
            [] => vec![None],
            [only] => vec![Some(*only)],
            [lhs, rhs] => {
                let sum = $builder.add_primitive(
                    $OpAdd,
                    vec![PrimitiveValue::Local(*lhs), PrimitiveValue::Local(*rhs)],
                    OperationRole::Linearized {
                        active_mask: vec![true, true],
                    },
                );
                vec![Some(sum[0])]
            }
            _ => unreachable!("mul linearization creates at most two terms"),
        })
    }};
}

#[macro_export]
macro_rules! linearize_exp {
    ($builder:expr, $OpMul:path, $primal_out:expr, $t0:expr) => {
        Ok(match $t0 {
            Some(dx) => {
                let out = $builder.add_primitive(
                    $OpMul,
                    vec![
                        PrimitiveValue::External($primal_out.clone()),
                        PrimitiveValue::Local(dx),
                    ],
                    OperationRole::Linearized {
                        active_mask: vec![false, true],
                    },
                );
                vec![Some(out[0])]
            }
            None => vec![None],
        })
    };
}

#[macro_export]
macro_rules! linearize_neg {
    ($builder:expr, $OpNeg:path, $t0:expr) => {
        Ok(match $t0 {
            Some(dx) => {
                let out = $builder.add_primitive(
                    $OpNeg,
                    vec![PrimitiveValue::Local(dx)],
                    OperationRole::Linearized {
                        active_mask: vec![true],
                    },
                );
                vec![Some(out[0])]
            }
            None => vec![None],
        })
    };
}

#[macro_export]
macro_rules! linearize_conj {
    ($builder:expr, $OpConj:path, $t0:expr) => {
        Ok(match $t0 {
            Some(dx) => {
                let out = $builder.add_primitive(
                    $OpConj,
                    vec![PrimitiveValue::Local(dx)],
                    OperationRole::Linearized {
                        active_mask: vec![true],
                    },
                );
                vec![Some(out[0])]
            }
            None => vec![None],
        })
    };
}