tidu/rules/
primitive_builder.rs

1use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey, ValueRef};
2
3/// Reference to a value available to a primitive AD rule.
4#[derive(Clone, Debug, PartialEq, Eq, Hash)]
5pub enum PrimitiveValue<Op: GraphOperation> {
6    /// Value produced inside the graph being built.
7    Local(LocalValueId),
8    /// Value from the source primitive computation graph.
9    External(computegraph::ValueKey<Op>),
10}
11
12/// Input reference passed to a primitive transpose rule.
13///
14/// Linearized graphs contain both residual values and values that belong to
15/// the tangent flow. A transpose rule may freely use residual values as
16/// ordinary graph operands, but tangent-flow values must not be retained as
17/// tensor-valued operands in the transposed graph. When a tangent-flow value is
18/// the tangent of a known primal value, `primal` carries that counterpart for
19/// metadata or runtime-shape use by downstream primitive sets.
20#[derive(Clone, Debug, PartialEq, Eq, Hash)]
21pub enum PrimitiveTransposeInput<Op: GraphOperation> {
22    /// Value that is independent of the differentiated tangent flow.
23    Residual(ValueKey<Op>),
24    /// Value on the differentiated tangent flow.
25    Linear {
26        /// Key of the linearized graph value.
27        key: ValueKey<Op>,
28        /// Primal value whose tangent this value represents, when known.
29        primal: Option<ValueKey<Op>>,
30    },
31}
32
33impl<Op: GraphOperation> PrimitiveTransposeInput<Op> {
34    /// Return the key of the linearized graph value.
35    pub fn key(&self) -> &ValueKey<Op> {
36        match self {
37            Self::Residual(key) | Self::Linear { key, .. } => key,
38        }
39    }
40
41    /// Return the corresponding primal key for a linear input, when known.
42    pub fn primal(&self) -> Option<&ValueKey<Op>> {
43        match self {
44            Self::Residual(_) => None,
45            Self::Linear { primal, .. } => primal.as_ref(),
46        }
47    }
48
49    /// Return this input as a value usable by rule builders when it is residual.
50    pub fn as_residual_value(&self) -> Option<PrimitiveValue<Op>> {
51        match self {
52            Self::Residual(key) => Some(PrimitiveValue::External(key.clone())),
53            Self::Linear { .. } => None,
54        }
55    }
56}
57
58impl<Op: GraphOperation> From<PrimitiveValue<Op>> for ValueRef<Op> {
59    fn from(value: PrimitiveValue<Op>) -> Self {
60        match value {
61            PrimitiveValue::Local(id) => ValueRef::Local(id),
62            PrimitiveValue::External(key) => ValueRef::External(key),
63        }
64    }
65}
66
67impl<Op: GraphOperation> From<ValueRef<Op>> for PrimitiveValue<Op> {
68    fn from(value: ValueRef<Op>) -> Self {
69        match value {
70            ValueRef::Local(id) => PrimitiveValue::Local(id),
71            ValueRef::External(key) => PrimitiveValue::External(key),
72        }
73    }
74}
75
76/// Builder used by primitive JVP and transpose rules to append primitive applications.
77pub trait PrimitiveBuilder<Op: GraphOperation> {
78    /// Add one primitive application and return local ids for its outputs.
79    fn add_primitive(
80        &mut self,
81        op: Op,
82        inputs: Vec<PrimitiveValue<Op>>,
83        role: OperationRole,
84    ) -> Vec<LocalValueId>;
85}
86
87#[cfg(test)]
88mod tests;
tidu/rules/primitive_builder.rs

tidu/rules/
primitive_builder.rs
