tidu/rules/
primitive_op.rs

1use super::{ADKey, ADRuleResult, PrimitiveBuilder, PrimitiveTransposeInput};
2use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey};
3
4/// Extends `GraphOperation` with primitive JVP and transpose rules for AD.
5///
6/// - `jvp_rule` is called by [`crate::linearize`]
7/// - `transpose_rule` is called by [`crate::linear_transpose`]
8///
9/// Both methods add new primitive applications through a [`PrimitiveBuilder`]. The downstream
10/// implementor is responsible for ensuring closure: every op emitted must also
11/// implement `Primitive`.
12///
13/// # Examples
14///
15/// ```
16/// use computegraph::{ValueKey, GraphOperation, LocalValueId, OperationRole};
17/// use tidu::{
18///     ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue,
19/// };
20///
21/// #[derive(Clone, Debug, PartialEq, Eq, Hash)]
22/// enum Key { Base(String), Tan(Box<Key>, DiffPassId) }
23///
24/// impl ADKey for Key {
25///     fn tangent_of(&self, p: DiffPassId) -> Self { Key::Tan(Box::new(self.clone()), p) }
26/// }
27///
28/// #[derive(Clone, Debug, PartialEq, Eq, Hash)]
29/// struct AddOp;
30///
31/// impl GraphOperation for AddOp {
32///     type Operand = f64;
33///     type Context = ();
34///     type InputKey = Key;
35///     fn input_count(&self) -> usize { 2 }
36///     fn output_count(&self) -> usize { 1 }
37/// }
38///
39/// impl Primitive for AddOp {
40///     type ADContext = ();
41///
42///     fn add() -> Self { AddOp }
43///     fn jvp_rule(
44///         &self, _b: &mut impl PrimitiveBuilder<Self>,
45///         _pi: &[ValueKey<Self>], _po: &[ValueKey<Self>],
46///         t: &[Option<LocalValueId>],
47///         _ctx: &mut (),
48///     ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
49///         Ok(vec![t[0].or(t[1])])
50///     }
51///     fn transpose_rule(
52///         &self, _builder: &mut impl PrimitiveBuilder<Self>,
53///         ct: &[Option<LocalValueId>], _i: &[PrimitiveTransposeInput<Self>], _m: &OperationRole,
54///         _ctx: &mut (),
55///     ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
56///         Ok(vec![ct[0], ct[0]])
57///     }
58/// }
59/// ```
60pub trait Primitive: GraphOperation
61where
62    Self::InputKey: ADKey,
63{
64    /// Runtime AD context threaded through linearization and transposition.
65    ///
66    /// This can carry information such as concrete shapes or guard decisions
67    /// that influence how AD rules emit graph structure.
68    type ADContext: Default;
69
70    /// Returns the addition operation used for cotangent accumulation
71    /// in [`crate::linear_transpose`]. When multiple cotangents flow to the same
72    /// `ValueKey`, `linear_transpose` emits `Op::add()` nodes to sum them.
73    fn add() -> Self
74    where
75        Self: Sized;
76
77    /// Emit the JVP rule for this primitive.
78    ///
79    /// Must be linear in tangent inputs. May reference primal inputs/outputs
80    /// through `External(ValueKey)`. Must emit ops in `OperationRole::Linearized`.
81    fn jvp_rule(
82        &self,
83        builder: &mut impl PrimitiveBuilder<Self>,
84        primal_inputs: &[ValueKey<Self>],
85        primal_outputs: &[ValueKey<Self>],
86        tangent_inputs: &[Option<LocalValueId>],
87        ctx: &mut Self::ADContext,
88    ) -> ADRuleResult<Vec<Option<LocalValueId>>>
89    where
90        Self: Sized;
91
92    /// Emit the transpose rule for this linear primitive.
93    ///
94    /// Receives cotangent outputs and produces cotangent inputs.
95    /// Must only emit ops that themselves implement `Primitive`.
96    fn transpose_rule(
97        &self,
98        builder: &mut impl PrimitiveBuilder<Self>,
99        cotangent_outputs: &[Option<LocalValueId>],
100        inputs: &[PrimitiveTransposeInput<Self>],
101        role: &OperationRole,
102        ctx: &mut Self::ADContext,
103    ) -> ADRuleResult<Vec<Option<LocalValueId>>>
104    where
105        Self: Sized;
106}
tidu/rules/primitive_op.rs

tidu/rules/
primitive_op.rs
