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 semantic rules are owned by
5//! `tenferro-ad`; the private dispatcher below is only the context-owned bridge
6//! used while the core sweep remains graph based.
7
8pub mod context;
9
10#[cfg(feature = "autodiff")]
11mod analytic;
12#[cfg(feature = "autodiff")]
13mod contraction;
14#[cfg(feature = "autodiff")]
15mod diagonal;
16#[cfg(feature = "autodiff")]
17mod dynamic;
18#[cfg(feature = "autodiff")]
19mod elementwise;
20#[cfg(feature = "autodiff")]
21mod indexing;
22#[cfg(feature = "autodiff")]
23pub(crate) mod registry;
24#[cfg(feature = "autodiff")]
25mod semiring;
26#[cfg(feature = "autodiff")]
27mod structural;
28#[cfg(feature = "autodiff")]
29#[doc(hidden)]
30pub mod support;
31#[cfg(feature = "autodiff")]
32#[doc(hidden)]
33pub mod transpose_input;
34#[cfg(feature = "autodiff")]
35mod zeros;
36
37#[cfg(feature = "autodiff")]
38use crate::ext_op::ExtensionOp;
39#[cfg(feature = "autodiff")]
40use crate::std_tensor_op::StdTensorOp;
41#[cfg(feature = "autodiff")]
42use computegraph::graph::GraphBuilder;
43#[cfg(feature = "autodiff")]
44use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef};
45
46#[cfg(feature = "autodiff")]
47#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48pub enum ADRuleKind {
49    Jvp,
50    Transpose,
51}
52
53#[cfg(feature = "autodiff")]
54#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
55pub enum ADRuleError {
56    #[error("unsupported AD rule {rule:?} for {op}")]
57    Unsupported { op: String, rule: ADRuleKind },
58    #[error("invalid AD rule input for {op} ({rule:?}): {message}")]
59    InvalidInput {
60        op: String,
61        rule: ADRuleKind,
62        message: String,
63    },
64}
65
66#[cfg(feature = "autodiff")]
67impl ADRuleError {
68    pub fn unsupported(op: impl Into<String>, rule: ADRuleKind) -> Self {
69        Self::Unsupported {
70            op: op.into(),
71            rule,
72        }
73    }
74
75    pub fn invalid_input(
76        op: impl Into<String>,
77        rule: ADRuleKind,
78        message: impl Into<String>,
79    ) -> Self {
80        Self::InvalidInput {
81            op: op.into(),
82            rule,
83            message: message.into(),
84        }
85    }
86
87    pub const fn rule(&self) -> ADRuleKind {
88        match self {
89            Self::Unsupported { rule, .. } | Self::InvalidInput { rule, .. } => *rule,
90        }
91    }
92}
93
94#[cfg(feature = "autodiff")]
95pub type ADRuleResult<T> = std::result::Result<T, ADRuleError>;
96
97#[cfg(feature = "autodiff")]
98#[derive(Clone, Debug)]
99pub enum PrimitiveTransposeInput<Op: computegraph::GraphOperation> {
100    Residual(ValueKey<Op>),
101    Linear {
102        key: ValueKey<Op>,
103        primal: Option<ValueKey<Op>>,
104    },
105}
106
107#[cfg(feature = "autodiff")]
108impl<Op: computegraph::GraphOperation> PrimitiveTransposeInput<Op> {
109    pub fn key(&self) -> &ValueKey<Op> {
110        match self {
111            Self::Residual(key) | Self::Linear { key, .. } => key,
112        }
113    }
114}
115
116#[cfg(feature = "autodiff")]
117fn missing_primitive_kind(op: &StdTensorOp, rule: ADRuleKind) -> ADRuleError {
118    ADRuleError::invalid_input(
119        "tenferro-internal-ops primitive AD dispatch",
120        rule,
121        format!("non-extension operation has no primitive kind: {op:?}"),
122    )
123}
124
125/// Builder interface used by tenferro AD rules.
126///
127/// # Examples
128///
129/// ```
130/// use computegraph::graph::GraphBuilder;
131/// use computegraph::{OperationRole, ValueRef};
132/// use tenferro_ops::ad::PrimitiveRuleBuilder;
133/// use tenferro_ops::input_key::TensorInputKey;
134/// use tenferro_ops::std_tensor_op::StdTensorOp;
135///
136/// let mut builder = GraphBuilder::<StdTensorOp>::new();
137/// let x = builder.add_input(TensorInputKey::User { id: 1 });
138/// let out = PrimitiveRuleBuilder::add_operation(
139///     &mut builder,
140///     StdTensorOp::Neg,
141///     vec![ValueRef::Local(x)],
142///     OperationRole::Primary,
143/// );
144/// assert_eq!(out.len(), 1);
145/// ```
146#[cfg(feature = "autodiff")]
147pub trait PrimitiveRuleBuilder {
148    /// Add one primitive graph operation and return local ids for its outputs.
149    fn add_operation(
150        &mut self,
151        operation: StdTensorOp,
152        inputs: Vec<ValueRef<StdTensorOp>>,
153        role: OperationRole,
154    ) -> Vec<LocalValueId>;
155}
156
157#[cfg(feature = "autodiff")]
158impl PrimitiveRuleBuilder for GraphBuilder<StdTensorOp> {
159    fn add_operation(
160        &mut self,
161        operation: StdTensorOp,
162        inputs: Vec<ValueRef<StdTensorOp>>,
163        role: OperationRole,
164    ) -> Vec<LocalValueId> {
165        GraphBuilder::add_operation(self, operation, inputs, role)
166    }
167}
168
169/// Single context-owned bridge from the core AD sweep to semantic extension
170/// rules.
171///
172/// Family-specific rules must not implement this trait. `tenferro-ad` owns the
173/// sole implementation and dispatches through its `SemanticExtensionRuleSet`.
174#[doc(hidden)]
175#[cfg(feature = "autodiff")]
176pub trait ExtensionAdDispatcher: std::fmt::Debug + Send + Sync + 'static {
177    fn has_primal_vjp(&self, family_id: &str) -> bool;
178
179    fn linearize(
180        &self,
181        op: &dyn ExtensionOp,
182        builder: &mut dyn PrimitiveRuleBuilder,
183        primal_in: &[ValueKey<StdTensorOp>],
184        primal_out: &[ValueKey<StdTensorOp>],
185        tangent_in: &[Option<LocalValueId>],
186        ctx: &mut context::ShapeGuardContext,
187    ) -> ADRuleResult<Vec<Option<LocalValueId>>>;
188
189    fn transpose(
190        &self,
191        op: &dyn ExtensionOp,
192        builder: &mut dyn PrimitiveRuleBuilder,
193        cotangent_out: &[Option<LocalValueId>],
194        inputs: &[PrimitiveTransposeInput<StdTensorOp>],
195        mode: &OperationRole,
196        ctx: &mut context::ShapeGuardContext,
197    ) -> ADRuleResult<Vec<Option<LocalValueId>>>;
198}
199
200#[cfg(feature = "autodiff")]
201fn dispatch_extension_linearize(
202    op: &dyn ExtensionOp,
203    builder: &mut dyn PrimitiveRuleBuilder,
204    primal_in: &[ValueKey<StdTensorOp>],
205    primal_out: &[ValueKey<StdTensorOp>],
206    tangent_in: &[Option<LocalValueId>],
207    ctx: &mut context::ShapeGuardContext,
208) -> ADRuleResult<Vec<Option<LocalValueId>>> {
209    if let Some(dispatcher) = ctx.extension_ad_dispatcher() {
210        return dispatcher.linearize(op, builder, primal_in, primal_out, tangent_in, ctx);
211    }
212    Err(ADRuleError::unsupported(op.family_id(), ADRuleKind::Jvp))
213}
214
215#[cfg(feature = "autodiff")]
216fn dispatch_extension_transpose(
217    op: &dyn ExtensionOp,
218    builder: &mut dyn PrimitiveRuleBuilder,
219    cotangent_out: &[Option<LocalValueId>],
220    inputs: &[PrimitiveTransposeInput<StdTensorOp>],
221    mode: &OperationRole,
222    ctx: &mut context::ShapeGuardContext,
223) -> ADRuleResult<Vec<Option<LocalValueId>>> {
224    if let Some(dispatcher) = ctx.extension_ad_dispatcher() {
225        return dispatcher.transpose(op, builder, cotangent_out, inputs, mode, ctx);
226    }
227    Err(ADRuleError::unsupported(
228        op.family_id(),
229        ADRuleKind::Transpose,
230    ))
231}
232
233/// Forward-mode AD (JVP) for `StdTensorOp`: given the primal op and its
234/// tangent inputs, emit the linearized graph into `builder` and return
235/// the output tangents.
236///
237/// Rules per op live in the category submodules (`semiring`, `analytic`,
238/// `elementwise`, `structural`, `contraction`, `indexing`, `diagonal`,
239/// `dynamic`). `StdTensorOp::Extension(_)` delegates to the trait.
240///
241/// # Errors
242///
243/// Returns [`ADRuleError::InvalidInput`] when the operation is not a known
244/// primitive or when a registered rule rejects the graph metadata. Returns
245/// [`ADRuleError::Unsupported`] when no AD rule is registered for the
246/// operation. Errors returned by an extension rule are propagated unchanged.
247#[cfg(feature = "autodiff")]
248pub fn linearize(
249    op: &StdTensorOp,
250    builder: &mut dyn PrimitiveRuleBuilder,
251    primal_in: &[ValueKey<StdTensorOp>],
252    primal_out: &[ValueKey<StdTensorOp>],
253    tangent_in: &[Option<LocalValueId>],
254    ctx: &mut context::ShapeGuardContext,
255) -> ADRuleResult<Vec<Option<LocalValueId>>> {
256    if let StdTensorOp::Extension(ext) = op {
257        return dispatch_extension_linearize(
258            ext.as_ref(),
259            builder,
260            primal_in,
261            primal_out,
262            tangent_in,
263            ctx,
264        );
265    }
266
267    let kind = op
268        .primitive_kind()
269        .ok_or_else(|| missing_primitive_kind(op, ADRuleKind::Jvp))?;
270    let rule = registry::primitive_ad_rule(kind)
271        .ok_or_else(|| registry::missing_rule(kind, ADRuleKind::Jvp))?;
272    rule.linearize(op, builder, primal_in, primal_out, tangent_in, ctx)
273}
274
275/// Reverse-mode AD (VJP) for `StdTensorOp`: given the primal op, its
276/// inputs, and the output cotangent, emit the transposed graph and
277/// return the input cotangents.
278///
279/// See [`linearize`] for the category split; the same categories appear
280/// here.
281///
282/// # Errors
283///
284/// Returns [`ADRuleError::InvalidInput`] when the operation, transpose inputs,
285/// or graph metadata are inconsistent. Returns [`ADRuleError::Unsupported`]
286/// when the required primitive or extension transpose rule is unavailable.
287/// Errors returned by a registered rule are propagated unchanged.
288#[cfg(feature = "autodiff")]
289pub fn transpose_rule(
290    op: &StdTensorOp,
291    builder: &mut impl PrimitiveRuleBuilder,
292    cotangent_out: &[Option<LocalValueId>],
293    inputs: &[PrimitiveTransposeInput<StdTensorOp>],
294    mode: &OperationRole,
295    ctx: &mut context::ShapeGuardContext,
296) -> ADRuleResult<Vec<Option<LocalValueId>>> {
297    if let StdTensorOp::Extension(ext) = op {
298        let builder_dyn: &mut dyn PrimitiveRuleBuilder = builder;
299        return dispatch_extension_transpose(
300            ext.as_ref(),
301            builder_dyn,
302            cotangent_out,
303            inputs,
304            mode,
305            ctx,
306        );
307    }
308
309    let transpose_inputs = inputs
310        .iter()
311        .map(transpose_input::TransposeInputRef::new)
312        .collect::<Vec<_>>();
313    let kind = op
314        .primitive_kind()
315        .ok_or_else(|| missing_primitive_kind(op, ADRuleKind::Transpose))?;
316    let rule = registry::primitive_ad_rule(kind)
317        .ok_or_else(|| registry::missing_rule(kind, ADRuleKind::Transpose))?;
318    let builder_dyn: &mut dyn PrimitiveRuleBuilder = builder;
319    rule.transpose_rule(op, builder_dyn, cotangent_out, &transpose_inputs, mode, ctx)
320}
321
322#[cfg(test)]
323mod tests;
tenferro_ops/ad/mod.rs

tenferro_ops/ad/
mod.rs
