tidu/
linearized_graph.rs

1use std::sync::Arc;
2
3use computegraph::graph::Graph;
4use computegraph::{GraphOperation, LocalValueId, ValueKey};
5
6/// Graph produced by linearizing a primitive computation graph.
7pub struct LinearizedGraph<Op: GraphOperation> {
8    linear: Graph<Op>,
9    residual: Arc<Graph<Op>>,
10    tangent_inputs: Vec<(Op::InputKey, LocalValueId)>,
11    tangent_outputs: Vec<Option<LocalValueId>>,
12    linear_primals: Vec<Option<ValueKey<Op>>>,
13}
14
15impl<Op: GraphOperation> LinearizedGraph<Op> {
16    pub(crate) fn from_parts(
17        linear: Graph<Op>,
18        residual: Arc<Graph<Op>>,
19        tangent_inputs: Vec<(Op::InputKey, LocalValueId)>,
20        tangent_outputs: Vec<Option<LocalValueId>>,
21        linear_primals: Vec<Option<ValueKey<Op>>>,
22    ) -> Self {
23        debug_assert_eq!(linear.values().len(), linear_primals.len());
24        Self {
25            linear,
26            residual,
27            tangent_inputs,
28            tangent_outputs,
29            linear_primals,
30        }
31    }
32
33    /// Borrow the strictly-linear graph representation.
34    pub fn as_graph(&self) -> &Graph<Op> {
35        &self.linear
36    }
37
38    /// Borrow the residual graph referenced by the linear graph.
39    pub fn residual_graph(&self) -> &Graph<Op> {
40        &self.residual
41    }
42
43    /// Consume this value and return the strictly-linear graph representation.
44    pub fn into_graph(self) -> Graph<Op> {
45        self.linear
46    }
47
48    /// Consume this value and return both graph parts.
49    pub fn into_graphs(self) -> (Graph<Op>, Arc<Graph<Op>>) {
50        (self.linear, self.residual)
51    }
52
53    /// Tangent input keys and local value ids.
54    pub fn tangent_inputs(&self) -> &[(Op::InputKey, LocalValueId)] {
55        &self.tangent_inputs
56    }
57
58    /// Tangent outputs aligned with requested primal outputs.
59    pub fn tangent_outputs(&self) -> &[Option<LocalValueId>] {
60        &self.tangent_outputs
61    }
62
63    pub(crate) fn linear_primal(&self, id: LocalValueId) -> Option<&ValueKey<Op>> {
64        self.linear_primals.get(id).and_then(Option::as_ref)
65    }
66}
tidu/linearized_graph.rs

tidu/
linearized_graph.rs
