tidu 0.1.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

use computegraph::graph::GraphBuilder;
use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueRef};

/// Reference to a value available to a primitive AD rule.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum PrimitiveValue<Op: GraphOperation> {
    /// Value produced inside the graph being built.
    Local(LocalValueId),
    /// Value from the source primitive computation graph.
    External(computegraph::ValueKey<Op>),
}

impl<Op: GraphOperation> From<PrimitiveValue<Op>> for ValueRef<Op> {
    fn from(value: PrimitiveValue<Op>) -> Self {
        match value {
            PrimitiveValue::Local(id) => ValueRef::Local(id),
            PrimitiveValue::External(key) => ValueRef::External(key),
        }
    }
}

impl<Op: GraphOperation> From<ValueRef<Op>> for PrimitiveValue<Op> {
    fn from(value: ValueRef<Op>) -> Self {
        match value {
            ValueRef::Local(id) => PrimitiveValue::Local(id),
            ValueRef::External(key) => PrimitiveValue::External(key),
        }
    }
}

/// Builder used by primitive JVP and transpose rules to append primitive applications.
pub trait PrimitiveBuilder<Op: GraphOperation> {
    /// Add one primitive application and return local ids for its outputs.
    fn add_primitive(
        &mut self,
        op: Op,
        inputs: Vec<PrimitiveValue<Op>>,
        role: OperationRole,
    ) -> Vec<LocalValueId>;
}

pub(crate) struct GraphPrimitiveBuilder<'a, Op: GraphOperation> {
    inner: &'a mut GraphBuilder<Op>,
}

impl<'a, Op: GraphOperation> GraphPrimitiveBuilder<'a, Op> {
    pub(crate) fn new(inner: &'a mut GraphBuilder<Op>) -> Self {
        Self { inner }
    }
}

impl<Op: GraphOperation> PrimitiveBuilder<Op> for GraphPrimitiveBuilder<'_, Op> {
    fn add_primitive(
        &mut self,
        op: Op,
        inputs: Vec<PrimitiveValue<Op>>,
        role: OperationRole,
    ) -> Vec<LocalValueId> {
        let inputs = inputs.into_iter().map(ValueRef::from).collect();
        self.inner.add_operation(op, inputs, role)
    }
}