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;
tidu/eager/trace.rs

tidu/eager/
trace.rs
