tidu 0.1.0

Automatic-differentiation transforms (linearize, transpose, eager reverse-mode) for primitive computation graphs.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

use super::{ADKey, ADRuleResult, PrimitiveBuilder, PrimitiveValue};
use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey};

/// Extends `GraphOperation` with primitive JVP and transpose rules for AD.
///
/// - `try_jvp_rule` is called by [`crate::try_linearize`]
/// - `try_linear_transpose_rule` is called by [`crate::try_linear_transpose`]
///
/// Both methods add new primitive applications through a [`PrimitiveBuilder`]. The downstream
/// implementor is responsible for ensuring closure: every op emitted must also
/// implement `Primitive`.
///
/// # Examples
///
/// ```
/// use computegraph::{ValueKey, GraphOperation, LocalValueId, OperationRole};
/// use tidu::{ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue};
///
/// #[derive(Clone, Debug, PartialEq, Eq, Hash)]
/// enum Key { Base(String), Tan(Box<Key>, DiffPassId) }
///
/// impl ADKey for Key {
///     fn tangent_of(&self, p: DiffPassId) -> Self { Key::Tan(Box::new(self.clone()), p) }
/// }
///
/// #[derive(Clone, Debug, PartialEq, Eq, Hash)]
/// struct AddOp;
///
/// impl GraphOperation for AddOp {
///     type Operand = f64;
///     type Context = ();
///     type InputKey = Key;
///     fn input_count(&self) -> usize { 2 }
///     fn output_count(&self) -> usize { 1 }
/// }
///
/// impl Primitive for AddOp {
///     type ADContext = ();
///
///     fn add() -> Self { AddOp }
///     fn jvp_rule(
///         &self, _b: &mut impl PrimitiveBuilder<Self>,
///         _pi: &[ValueKey<Self>], _po: &[ValueKey<Self>],
///         t: &[Option<LocalValueId>],
///         _ctx: &mut (),
///     ) -> Vec<Option<LocalValueId>> {
///         vec![t[0].or(t[1])]
///     }
///     fn transpose_rule(
///         &self, _builder: &mut impl PrimitiveBuilder<Self>,
///         ct: &[Option<LocalValueId>], _i: &[PrimitiveValue<Self>], _m: &OperationRole,
///         _ctx: &mut (),
///     ) -> Vec<Option<LocalValueId>> {
///         vec![ct[0], ct[0]]
///     }
/// }
/// ```
pub trait Primitive: GraphOperation
where
    Self::InputKey: ADKey,
{
    /// Runtime AD context threaded through linearization and transposition.
    ///
    /// This can carry information such as concrete shapes or guard decisions
    /// that influence how AD rules emit graph structure.
    type ADContext: Default;

    /// Returns the addition operation used for cotangent accumulation
    /// in [`crate::linear_transpose`]. When multiple cotangents flow to the same
    /// `ValueKey`, `linear_transpose` emits `Op::add()` nodes to sum them.
    fn add() -> Self
    where
        Self: Sized;

    /// Emit the JVP rule for this primitive.
    ///
    /// Must be linear in tangent inputs. May reference primal inputs/outputs
    /// through `External(ValueKey)`. Must emit ops in `OperationRole::Linearized`.
    fn jvp_rule(
        &self,
        builder: &mut impl PrimitiveBuilder<Self>,
        primal_inputs: &[ValueKey<Self>],
        primal_outputs: &[ValueKey<Self>],
        tangent_inputs: &[Option<LocalValueId>],
        ctx: &mut Self::ADContext,
    ) -> Vec<Option<LocalValueId>>
    where
        Self: Sized;

    /// Fallible variant of [`Primitive::jvp_rule`].
    ///
    /// Implementors that can encounter missing extension rules should override
    /// this method and return [`super::ADRuleError`] instead of panicking. The
    /// default implementation preserves the infallible contract for existing
    /// primitive sets.
    fn try_jvp_rule(
        &self,
        builder: &mut impl PrimitiveBuilder<Self>,
        primal_inputs: &[ValueKey<Self>],
        primal_outputs: &[ValueKey<Self>],
        tangent_inputs: &[Option<LocalValueId>],
        ctx: &mut Self::ADContext,
    ) -> ADRuleResult<Vec<Option<LocalValueId>>>
    where
        Self: Sized,
    {
        Ok(self.jvp_rule(builder, primal_inputs, primal_outputs, tangent_inputs, ctx))
    }

    /// Emit the transpose rule for this linear primitive.
    ///
    /// Receives cotangent outputs and produces cotangent inputs.
    /// Must only emit ops that themselves implement `Primitive`.
    fn transpose_rule(
        &self,
        builder: &mut impl PrimitiveBuilder<Self>,
        cotangent_outputs: &[Option<LocalValueId>],
        inputs: &[PrimitiveValue<Self>],
        role: &OperationRole,
        ctx: &mut Self::ADContext,
    ) -> Vec<Option<LocalValueId>>
    where
        Self: Sized;

    /// Fallible variant of [`Primitive::transpose_rule`].
    ///
    /// Implementors that can encounter missing extension rules should override
    /// this method and return [`super::ADRuleError`] instead of panicking. The
    /// default implementation preserves the infallible contract for existing
    /// primitive sets.
    fn try_linear_transpose_rule(
        &self,
        builder: &mut impl PrimitiveBuilder<Self>,
        cotangent_outputs: &[Option<LocalValueId>],
        inputs: &[PrimitiveValue<Self>],
        role: &OperationRole,
        ctx: &mut Self::ADContext,
    ) -> ADRuleResult<Vec<Option<LocalValueId>>>
    where
        Self: Sized,
    {
        Ok(self.transpose_rule(builder, cotangent_outputs, inputs, role, ctx))
    }
}