tenferro_ops/ad/
mod.rs

1//! Automatic differentiation rules for [`StdTensorOp`].
2//!
3//! `linearize` and `transpose_rule` are separate graph-level contracts.
4//! Core ops keep their rules here; extension ops own their own AD support
5//! through the extension trait.
6
7pub mod context;
8
9#[cfg(feature = "autodiff")]
10mod analytic;
11#[cfg(feature = "autodiff")]
12mod contraction;
13#[cfg(feature = "autodiff")]
14mod diagonal;
15#[cfg(feature = "autodiff")]
16mod dynamic;
17#[cfg(feature = "autodiff")]
18mod elementwise;
19#[cfg(feature = "autodiff")]
20mod indexing;
21#[cfg(feature = "autodiff")]
22pub(crate) mod registry;
23#[cfg(feature = "autodiff")]
24mod semiring;
25#[cfg(feature = "autodiff")]
26mod structural;
27#[cfg(feature = "autodiff")]
28#[doc(hidden)]
29pub mod support;
30#[cfg(feature = "autodiff")]
31#[doc(hidden)]
32pub mod transpose_input;
33#[cfg(feature = "autodiff")]
34mod zeros;
35
36#[cfg(feature = "autodiff")]
37use computegraph::graph::GraphBuilder;
38#[cfg(feature = "autodiff")]
39use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef};
40#[cfg(feature = "autodiff")]
41use tidu::{
42    ADRuleError, ADRuleKind, ADRuleResult, PrimitiveBuilder, PrimitiveTransposeInput,
43    PrimitiveValue,
44};
45
46#[cfg(feature = "autodiff")]
47use crate::ext_op::{linearize_extension_rule, transpose_extension_rule};
48#[cfg(feature = "autodiff")]
49use crate::std_tensor_op::StdTensorOp;
50
51#[cfg(feature = "autodiff")]
52fn missing_primitive_kind(op: &StdTensorOp, rule: ADRuleKind) -> ADRuleError {
53    ADRuleError::invalid_input(
54        "tenferro-internal-ops primitive AD dispatch",
55        rule,
56        format!("non-extension operation has no primitive kind: {op:?}"),
57    )
58}
59
60/// Builder interface used by tenferro AD rules.
61///
62/// # Examples
63///
64/// ```
65/// use computegraph::graph::GraphBuilder;
66/// use computegraph::{OperationRole, ValueRef};
67/// use tenferro_ops::ad::PrimitiveRuleBuilder;
68/// use tenferro_ops::input_key::TensorInputKey;
69/// use tenferro_ops::std_tensor_op::StdTensorOp;
70///
71/// let mut builder = GraphBuilder::<StdTensorOp>::new();
72/// let x = builder.add_input(TensorInputKey::User { id: 1 });
73/// let out = PrimitiveRuleBuilder::add_operation(
74///     &mut builder,
75///     StdTensorOp::Neg,
76///     vec![ValueRef::Local(x)],
77///     OperationRole::Primary,
78/// );
79/// assert_eq!(out.len(), 1);
80/// ```
81#[cfg(feature = "autodiff")]
82pub trait PrimitiveRuleBuilder {
83    /// Add one primitive graph operation and return local ids for its outputs.
84    fn add_operation(
85        &mut self,
86        operation: StdTensorOp,
87        inputs: Vec<ValueRef<StdTensorOp>>,
88        role: OperationRole,
89    ) -> Vec<LocalValueId>;
90}
91
92#[cfg(feature = "autodiff")]
93impl<B> PrimitiveRuleBuilder for B
94where
95    B: PrimitiveBuilder<StdTensorOp> + ?Sized,
96{
97    fn add_operation(
98        &mut self,
99        operation: StdTensorOp,
100        inputs: Vec<ValueRef<StdTensorOp>>,
101        role: OperationRole,
102    ) -> Vec<LocalValueId> {
103        let inputs = inputs.into_iter().map(PrimitiveValue::from).collect();
104        PrimitiveBuilder::add_primitive(self, operation, inputs, role)
105    }
106}
107
108#[cfg(feature = "autodiff")]
109impl PrimitiveRuleBuilder for GraphBuilder<StdTensorOp> {
110    fn add_operation(
111        &mut self,
112        operation: StdTensorOp,
113        inputs: Vec<ValueRef<StdTensorOp>>,
114        role: OperationRole,
115    ) -> Vec<LocalValueId> {
116        GraphBuilder::add_operation(self, operation, inputs, role)
117    }
118}
119
120/// Forward-mode AD (JVP) for `StdTensorOp`: given the primal op and its
121/// tangent inputs, emit the linearized graph into `builder` and return
122/// the output tangents.
123///
124/// Rules per op live in the category submodules (`semiring`, `analytic`,
125/// `elementwise`, `structural`, `contraction`, `indexing`, `diagonal`,
126/// `dynamic`). `StdTensorOp::Extension(_)` delegates to the trait.
127#[cfg(feature = "autodiff")]
128pub fn linearize(
129    op: &StdTensorOp,
130    builder: &mut dyn PrimitiveRuleBuilder,
131    primal_in: &[ValueKey<StdTensorOp>],
132    primal_out: &[ValueKey<StdTensorOp>],
133    tangent_in: &[Option<LocalValueId>],
134    ctx: &mut context::ShapeGuardContext,
135) -> ADRuleResult<Vec<Option<LocalValueId>>> {
136    if let StdTensorOp::Extension(ext) = op {
137        return linearize_extension_rule(
138            ext.as_ref(),
139            builder,
140            primal_in,
141            primal_out,
142            tangent_in,
143            ctx,
144        );
145    }
146
147    let kind = op
148        .primitive_kind()
149        .ok_or_else(|| missing_primitive_kind(op, ADRuleKind::Jvp))?;
150    let rule = registry::primitive_ad_rule(kind)
151        .ok_or_else(|| registry::missing_rule(kind, ADRuleKind::Jvp))?;
152    rule.linearize(op, builder, primal_in, primal_out, tangent_in, ctx)
153}
154
155/// Reverse-mode AD (VJP) for `StdTensorOp`: given the primal op, its
156/// inputs, and the output cotangent, emit the transposed graph and
157/// return the input cotangents.
158///
159/// See [`linearize`] for the category split; the same categories appear
160/// here.
161#[cfg(feature = "autodiff")]
162pub fn transpose_rule(
163    op: &StdTensorOp,
164    builder: &mut impl PrimitiveRuleBuilder,
165    cotangent_out: &[Option<LocalValueId>],
166    inputs: &[PrimitiveTransposeInput<StdTensorOp>],
167    mode: &OperationRole,
168    ctx: &mut context::ShapeGuardContext,
169) -> ADRuleResult<Vec<Option<LocalValueId>>> {
170    if let StdTensorOp::Extension(ext) = op {
171        let builder_dyn: &mut dyn PrimitiveRuleBuilder = builder;
172        return transpose_extension_rule(
173            ext.as_ref(),
174            builder_dyn,
175            cotangent_out,
176            inputs,
177            mode,
178            ctx,
179        );
180    }
181
182    let transpose_inputs = inputs
183        .iter()
184        .map(transpose_input::TransposeInputRef::new)
185        .collect::<Vec<_>>();
186    let kind = op
187        .primitive_kind()
188        .ok_or_else(|| missing_primitive_kind(op, ADRuleKind::Transpose))?;
189    let rule = registry::primitive_ad_rule(kind)
190        .ok_or_else(|| registry::missing_rule(kind, ADRuleKind::Transpose))?;
191    let builder_dyn: &mut dyn PrimitiveRuleBuilder = builder;
192    rule.transpose_rule(op, builder_dyn, cotangent_out, &transpose_inputs, mode, ctx)
193}
194
195#[cfg(test)]
196mod tests;
tenferro_ops/ad/mod.rs

tenferro_ops/ad/
mod.rs
