tidu/eager/
trace.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use computegraph::{GraphOperation, ValueKey};
5
6use super::record::{EagerRecordError, EagerRecordResult, RecordedGraph};
7
8/// Opaque handle to an eager reverse-mode trace node.
9///
10/// Downstream eager values store this handle next to their concrete data and
11/// pass it back to [`crate::eager::backward`]. The node and edge layout is
12/// intentionally private.
13pub struct Trace<Op: GraphOperation> {
14    node: Arc<TraceNode<Op>>,
15}
16
17impl<Op: GraphOperation> Clone for Trace<Op> {
18    fn clone(&self) -> Self {
19        Self {
20            node: self.node.clone(),
21        }
22    }
23}
24
25impl<Op: GraphOperation> Trace<Op> {
26    pub(crate) fn new(node: Arc<TraceNode<Op>>) -> Self {
27        Self { node }
28    }
29
30    pub(crate) fn node(&self) -> &Arc<TraceNode<Op>> {
31        &self.node
32    }
33
34    /// Saved concrete primal values used as initial data during backward replay.
35    pub fn saved_values(&self) -> &HashMap<ValueKey<Op>, Arc<Op::Operand>> {
36        self.node.saved_data()
37    }
38}
39
40pub(crate) struct TraceNode<Op: GraphOperation> {
41    computation: RecordedGraph<Op>,
42    primal_out_keys: Vec<ValueKey<Op>>,
43    saved_data: HashMap<ValueKey<Op>, Arc<Op::Operand>>,
44    input_edges: Vec<TraceEdge<Op>>,
45}
46
47impl<Op: GraphOperation> TraceNode<Op> {
48    pub(crate) fn new(
49        computation: RecordedGraph<Op>,
50        primal_out_keys: Vec<ValueKey<Op>>,
51        saved_data: HashMap<ValueKey<Op>, Arc<Op::Operand>>,
52        input_edges: Vec<TraceEdge<Op>>,
53    ) -> EagerRecordResult<Self> {
54        if primal_out_keys.len() != computation.output_keys().len() {
55            return Err(EagerRecordError::count_mismatch(
56                "TraceNode primal output keys",
57                computation.output_keys().len(),
58                primal_out_keys.len(),
59            ));
60        }
61        if input_edges.len() != computation.input_keys().len() {
62            return Err(EagerRecordError::count_mismatch(
63                "TraceNode input edges",
64                computation.input_keys().len(),
65                input_edges.len(),
66            ));
67        }
68
69        Ok(Self {
70            computation,
71            primal_out_keys,
72            saved_data,
73            input_edges,
74        })
75    }
76
77    pub(crate) fn computation(&self) -> &RecordedGraph<Op> {
78        &self.computation
79    }
80
81    pub(crate) fn primal_out_keys(&self) -> &[ValueKey<Op>] {
82        &self.primal_out_keys
83    }
84
85    pub(crate) fn saved_data(&self) -> &HashMap<ValueKey<Op>, Arc<Op::Operand>> {
86        &self.saved_data
87    }
88
89    pub(crate) fn input_edges(&self) -> &[TraceEdge<Op>] {
90        &self.input_edges
91    }
92}
93
94pub(crate) struct TraceEdge<Op: GraphOperation> {
95    pub(crate) node: Option<Arc<TraceNode<Op>>>,
96    pub(crate) key: ValueKey<Op>,
97    pub(crate) requires_grad: bool,
98}
99
100impl<Op: GraphOperation> TraceEdge<Op> {
101    pub(crate) fn new(
102        node: Option<Arc<TraceNode<Op>>>,
103        key: ValueKey<Op>,
104        requires_grad: bool,
105    ) -> Self {
106        Self {
107            node,
108            key,
109            requires_grad,
110        }
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use computegraph::GraphOperation;
118
119    #[derive(Clone, Debug, Hash, PartialEq, Eq)]
120    enum Op {
121        Id,
122    }
123
124    impl GraphOperation for Op {
125        type Operand = f64;
126        type Context = ();
127        type InputKey = &'static str;
128
129        fn input_count(&self) -> usize {
130            1
131        }
132
133        fn output_count(&self) -> usize {
134            1
135        }
136    }
137
138    fn recorded_id() -> RecordedGraph<Op> {
139        RecordedGraph::from_primitive(Op::Id, vec!["x"])
140            .expect("test primitive graph metadata should be valid")
141    }
142
143    #[test]
144    fn trace_node_rejects_mismatched_output_keys() {
145        let err = match TraceNode::new(
146            recorded_id(),
147            Vec::new(),
148            HashMap::new(),
149            vec![TraceEdge::new(None, ValueKey::Input("x"), true)],
150        ) {
151            Ok(_) => panic!("TraceNode::new unexpectedly accepted mismatched outputs"),
152            Err(err) => err,
153        };
154
155        assert!(matches!(
156            err,
157            EagerRecordError::CountMismatch {
158                field: "TraceNode primal output keys",
159                expected: 1,
160                actual: 0
161            }
162        ));
163    }
164
165    #[test]
166    fn trace_node_rejects_mismatched_input_edges() {
167        let output_key = recorded_id().output_keys()[0].clone();
168        let err = match TraceNode::new(recorded_id(), vec![output_key], HashMap::new(), Vec::new())
169        {
170            Ok(_) => panic!("TraceNode::new unexpectedly accepted mismatched input edges"),
171            Err(err) => err,
172        };
173
174        assert!(matches!(
175            err,
176            EagerRecordError::CountMismatch {
177                field: "TraceNode input edges",
178                expected: 1,
179                actual: 0
180            }
181        ));
182    }
183}
tidu/eager/trace.rs

tidu/eager/
trace.rs
