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

use std::collections::HashMap;
use std::sync::Arc;

use computegraph::{GraphOperation, ValueKey};

use super::record::{EagerRecordError, EagerRecordResult, RecordedGraph};

/// Opaque handle to an eager reverse-mode trace node.
///
/// Downstream eager values store this handle next to their concrete data and
/// pass it back to [`crate::eager::backward`]. The node and edge layout is
/// intentionally private.
pub struct Trace<Op: GraphOperation> {
    node: Arc<TraceNode<Op>>,
}

impl<Op: GraphOperation> Clone for Trace<Op> {
    fn clone(&self) -> Self {
        Self {
            node: self.node.clone(),
        }
    }
}

impl<Op: GraphOperation> Trace<Op> {
    pub(crate) fn new(node: Arc<TraceNode<Op>>) -> Self {
        Self { node }
    }

    pub(crate) fn node(&self) -> &Arc<TraceNode<Op>> {
        &self.node
    }

    /// Saved concrete primal values used as initial data during backward replay.
    pub fn saved_values(&self) -> &HashMap<ValueKey<Op>, Arc<Op::Operand>> {
        self.node.saved_data()
    }
}

pub(crate) struct TraceNode<Op: GraphOperation> {
    computation: RecordedGraph<Op>,
    primal_out_keys: Vec<ValueKey<Op>>,
    saved_data: HashMap<ValueKey<Op>, Arc<Op::Operand>>,
    input_edges: Vec<TraceEdge<Op>>,
}

impl<Op: GraphOperation> TraceNode<Op> {
    pub(crate) fn new(
        computation: RecordedGraph<Op>,
        primal_out_keys: Vec<ValueKey<Op>>,
        saved_data: HashMap<ValueKey<Op>, Arc<Op::Operand>>,
        input_edges: Vec<TraceEdge<Op>>,
    ) -> EagerRecordResult<Self> {
        if primal_out_keys.len() != computation.output_keys().len() {
            return Err(EagerRecordError::count_mismatch(
                "TraceNode primal output keys",
                computation.output_keys().len(),
                primal_out_keys.len(),
            ));
        }
        if input_edges.len() != computation.input_keys().len() {
            return Err(EagerRecordError::count_mismatch(
                "TraceNode input edges",
                computation.input_keys().len(),
                input_edges.len(),
            ));
        }

        Ok(Self {
            computation,
            primal_out_keys,
            saved_data,
            input_edges,
        })
    }

    pub(crate) fn computation(&self) -> &RecordedGraph<Op> {
        &self.computation
    }

    pub(crate) fn primal_out_keys(&self) -> &[ValueKey<Op>] {
        &self.primal_out_keys
    }

    pub(crate) fn saved_data(&self) -> &HashMap<ValueKey<Op>, Arc<Op::Operand>> {
        &self.saved_data
    }

    pub(crate) fn input_edges(&self) -> &[TraceEdge<Op>] {
        &self.input_edges
    }
}

pub(crate) struct TraceEdge<Op: GraphOperation> {
    pub(crate) node: Option<Arc<TraceNode<Op>>>,
    pub(crate) key: ValueKey<Op>,
    pub(crate) requires_grad: bool,
}

impl<Op: GraphOperation> TraceEdge<Op> {
    pub(crate) fn new(
        node: Option<Arc<TraceNode<Op>>>,
        key: ValueKey<Op>,
        requires_grad: bool,
    ) -> Self {
        Self {
            node,
            key,
            requires_grad,
        }
    }
}