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/// Per-rule declaration of which primal inputs/outputs a transpose (VJP) rule
98/// reads as full tensor residuals versus which need only shape/dtype metadata.
99///
100/// The AD engine uses this to bound residual retention: an index not declared
101/// here may only be accessed through its metadata, never as a tensor operand.
102/// Indices are counted in primal-input order (input mask) and primal-output
103/// order (output mask).
104///
105/// # Examples
106///
107/// ```
108/// use tenferro_ops::ad::ResidualSpec;
109///
110/// // `add` needs no tensor residuals; `mul` keeps both operands.
111/// let add = ResidualSpec::none();
112/// let mul = ResidualSpec::input(0).with_input(1);
113/// assert!(add.is_empty());
114/// assert!(mul.declares_input(0) && mul.declares_input(1));
115/// // Unary ops that reuse the forward output (e.g. exp) declare it.
116/// let exp = ResidualSpec::output(0);
117/// assert!(exp.declares_output(0));
118/// ```
119#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
120pub struct ResidualSpec {
121    input_mask: u64,
122    output_mask: u64,
123}
124
125impl ResidualSpec {
126    /// A mask declaring no tensor residuals (metadata-only rule).
127    pub const fn none() -> Self {
128        Self {
129            input_mask: 0,
130            output_mask: 0,
131        }
132    }
133
134    /// A mask declaring one input index as a tensor residual.
135    pub const fn input(index: usize) -> Self {
136        Self {
137            input_mask: 1 << index,
138            output_mask: 0,
139        }
140    }
141
142    /// A mask declaring one output index as a tensor residual.
143    pub const fn output(index: usize) -> Self {
144        Self {
145            input_mask: 0,
146            output_mask: 1 << index,
147        }
148    }
149
150    /// Add one input index to this mask.
151    pub const fn with_input(mut self, index: usize) -> Self {
152        self.input_mask |= 1 << index;
153        self
154    }
155
156    /// Add one output index to this mask.
157    pub const fn with_output(mut self, index: usize) -> Self {
158        self.output_mask |= 1 << index;
159        self
160    }
161
162    /// Add every input index to this mask.
163    pub const fn with_all_inputs(mut self) -> Self {
164        self.input_mask = u64::MAX;
165        self
166    }
167
168    /// Add every output index to this mask.
169    pub const fn with_all_outputs(mut self) -> Self {
170        self.output_mask = u64::MAX;
171        self
172    }
173
174    /// A mask declaring every input index as a tensor residual.
175    ///
176    /// Used by rules whose required operand set depends on the active-input
177    /// configuration (e.g. `mul`, concatenate, einsum), and by multi-op
178    /// families whose individual ops collectively read any operand.
179    pub const fn all_inputs() -> Self {
180        Self {
181            input_mask: u64::MAX,
182            output_mask: 0,
183        }
184    }
185
186    /// A mask declaring every output index as a tensor residual.
187    pub const fn all_outputs() -> Self {
188        Self {
189            input_mask: 0,
190            output_mask: u64::MAX,
191        }
192    }
193
194    /// Whether input `index` is declared as a tensor residual.
195    pub const fn declares_input(&self, index: usize) -> bool {
196        index < 64 && (self.input_mask >> index) & 1 == 1
197    }
198
199    /// Whether output `index` is declared as a tensor residual.
200    pub const fn declares_output(&self, index: usize) -> bool {
201        index < 64 && (self.output_mask >> index) & 1 == 1
202    }
203
204    /// Whether this mask declares no tensor residuals.
205    #[must_use]
206    pub const fn is_empty(&self) -> bool {
207        self.input_mask == 0 && self.output_mask == 0
208    }
209}
210
211#[cfg(feature = "autodiff")]
212#[derive(Clone, Debug)]
213pub enum PrimitiveTransposeInput<Op: computegraph::GraphOperation> {
214    Residual(ValueKey<Op>),
215    Linear {
216        key: ValueKey<Op>,
217        primal: Option<ValueKey<Op>>,
218    },
219}
220
221#[cfg(feature = "autodiff")]
222impl<Op: computegraph::GraphOperation> PrimitiveTransposeInput<Op> {
223    pub fn key(&self) -> &ValueKey<Op> {
224        match self {
225            Self::Residual(key) | Self::Linear { key, .. } => key,
226        }
227    }
228}
229
230#[cfg(feature = "autodiff")]
231fn missing_primitive_kind(op: &StdTensorOp, rule: ADRuleKind) -> ADRuleError {
232    ADRuleError::invalid_input(
233        "tenferro-internal-ops primitive AD dispatch",
234        rule,
235        format!("non-extension operation has no primitive kind: {op:?}"),
236    )
237}
238
239/// Builder interface used by tenferro AD rules.
240///
241/// # Examples
242///
243/// ```
244/// use computegraph::graph::GraphBuilder;
245/// use computegraph::{OperationRole, ValueRef};
246/// use tenferro_ops::ad::PrimitiveRuleBuilder;
247/// use tenferro_ops::input_key::TensorInputKey;
248/// use tenferro_ops::std_tensor_op::StdTensorOp;
249///
250/// let mut builder = GraphBuilder::<StdTensorOp>::new();
251/// let x = builder.add_input(TensorInputKey::User { id: 1 });
252/// let out = PrimitiveRuleBuilder::add_operation(
253///     &mut builder,
254///     StdTensorOp::Neg,
255///     vec![ValueRef::Local(x)],
256///     OperationRole::Primary,
257/// );
258/// assert_eq!(out.len(), 1);
259/// ```
260#[cfg(feature = "autodiff")]
261pub trait PrimitiveRuleBuilder {
262    /// Add one primitive graph operation and return local ids for its outputs.
263    fn add_operation(
264        &mut self,
265        operation: StdTensorOp,
266        inputs: Vec<ValueRef<StdTensorOp>>,
267        role: OperationRole,
268    ) -> Vec<LocalValueId>;
269}
270
271#[cfg(feature = "autodiff")]
272impl PrimitiveRuleBuilder for GraphBuilder<StdTensorOp> {
273    fn add_operation(
274        &mut self,
275        operation: StdTensorOp,
276        inputs: Vec<ValueRef<StdTensorOp>>,
277        role: OperationRole,
278    ) -> Vec<LocalValueId> {
279        GraphBuilder::add_operation(self, operation, inputs, role)
280    }
281}
282
283/// Single context-owned bridge from the core AD sweep to semantic extension
284/// rules.
285///
286/// Family-specific rules must not implement this trait. `tenferro-ad` owns the
287/// sole implementation and dispatches through its `SemanticExtensionRuleSet`.
288#[doc(hidden)]
289#[cfg(feature = "autodiff")]
290pub trait ExtensionAdDispatcher: std::fmt::Debug + Send + Sync + 'static {
291    fn has_primal_vjp(&self, family_id: &str) -> bool;
292
293    fn linearize(
294        &self,
295        op: &dyn ExtensionOp,
296        builder: &mut dyn PrimitiveRuleBuilder,
297        primal_in: &[ValueKey<StdTensorOp>],
298        primal_out: &[ValueKey<StdTensorOp>],
299        tangent_in: &[Option<LocalValueId>],
300        ctx: &mut context::ShapeGuardContext,
301    ) -> ADRuleResult<Vec<Option<LocalValueId>>>;
302
303    fn transpose(
304        &self,
305        op: &dyn ExtensionOp,
306        builder: &mut dyn PrimitiveRuleBuilder,
307        cotangent_out: &[Option<LocalValueId>],
308        inputs: &[PrimitiveTransposeInput<StdTensorOp>],
309        mode: &OperationRole,
310        ctx: &mut context::ShapeGuardContext,
311    ) -> ADRuleResult<Vec<Option<LocalValueId>>>;
312}
313
314#[cfg(feature = "autodiff")]
315fn dispatch_extension_linearize(
316    op: &dyn ExtensionOp,
317    builder: &mut dyn PrimitiveRuleBuilder,
318    primal_in: &[ValueKey<StdTensorOp>],
319    primal_out: &[ValueKey<StdTensorOp>],
320    tangent_in: &[Option<LocalValueId>],
321    ctx: &mut context::ShapeGuardContext,
322) -> ADRuleResult<Vec<Option<LocalValueId>>> {
323    if let Some(dispatcher) = ctx.extension_ad_dispatcher() {
324        return dispatcher.linearize(op, builder, primal_in, primal_out, tangent_in, ctx);
325    }
326    Err(ADRuleError::unsupported(op.family_id(), ADRuleKind::Jvp))
327}
328
329#[cfg(feature = "autodiff")]
330fn dispatch_extension_transpose(
331    op: &dyn ExtensionOp,
332    builder: &mut dyn PrimitiveRuleBuilder,
333    cotangent_out: &[Option<LocalValueId>],
334    inputs: &[PrimitiveTransposeInput<StdTensorOp>],
335    mode: &OperationRole,
336    ctx: &mut context::ShapeGuardContext,
337) -> ADRuleResult<Vec<Option<LocalValueId>>> {
338    if let Some(dispatcher) = ctx.extension_ad_dispatcher() {
339        return dispatcher.transpose(op, builder, cotangent_out, inputs, mode, ctx);
340    }
341    Err(ADRuleError::unsupported(
342        op.family_id(),
343        ADRuleKind::Transpose,
344    ))
345}
346
347/// Forward-mode AD (JVP) for `StdTensorOp`: given the primal op and its
348/// tangent inputs, emit the linearized graph into `builder` and return
349/// the output tangents.
350///
351/// Rules per op live in the category submodules (`semiring`, `analytic`,
352/// `elementwise`, `structural`, `contraction`, `indexing`, `diagonal`,
353/// `dynamic`). `StdTensorOp::Extension(_)` delegates to the trait.
354///
355/// # Errors
356///
357/// Returns [`ADRuleError::InvalidInput`] when the operation is not a known
358/// primitive or when a registered rule rejects the graph metadata. Returns
359/// [`ADRuleError::Unsupported`] when no AD rule is registered for the
360/// operation. Errors returned by an extension rule are propagated unchanged.
361#[cfg(feature = "autodiff")]
362pub fn linearize(
363    op: &StdTensorOp,
364    builder: &mut dyn PrimitiveRuleBuilder,
365    primal_in: &[ValueKey<StdTensorOp>],
366    primal_out: &[ValueKey<StdTensorOp>],
367    tangent_in: &[Option<LocalValueId>],
368    ctx: &mut context::ShapeGuardContext,
369) -> ADRuleResult<Vec<Option<LocalValueId>>> {
370    if let StdTensorOp::Extension(ext) = op {
371        return dispatch_extension_linearize(
372            ext.as_ref(),
373            builder,
374            primal_in,
375            primal_out,
376            tangent_in,
377            ctx,
378        );
379    }
380
381    let kind = op
382        .primitive_kind()
383        .ok_or_else(|| missing_primitive_kind(op, ADRuleKind::Jvp))?;
384    let rule = registry::primitive_ad_rule(kind)
385        .ok_or_else(|| registry::missing_rule(kind, ADRuleKind::Jvp))?;
386    rule.linearize(op, builder, primal_in, primal_out, tangent_in, ctx)
387}
388
389/// Reverse-mode AD (VJP) for `StdTensorOp`: given the primal op, its
390/// inputs, and the output cotangent, emit the transposed graph and
391/// return the input cotangents.
392///
393/// See [`linearize`] for the category split; the same categories appear
394/// here.
395///
396/// # Errors
397///
398/// Returns [`ADRuleError::InvalidInput`] when the operation, transpose inputs,
399/// or graph metadata are inconsistent. Returns [`ADRuleError::Unsupported`]
400/// when the required primitive or extension transpose rule is unavailable.
401/// Errors returned by a registered rule are propagated unchanged.
402#[cfg(feature = "autodiff")]
403pub fn transpose_rule(
404    op: &StdTensorOp,
405    builder: &mut impl PrimitiveRuleBuilder,
406    cotangent_out: &[Option<LocalValueId>],
407    inputs: &[PrimitiveTransposeInput<StdTensorOp>],
408    mode: &OperationRole,
409    ctx: &mut context::ShapeGuardContext,
410) -> ADRuleResult<Vec<Option<LocalValueId>>> {
411    if let StdTensorOp::Extension(ext) = op {
412        let builder_dyn: &mut dyn PrimitiveRuleBuilder = builder;
413        return dispatch_extension_transpose(
414            ext.as_ref(),
415            builder_dyn,
416            cotangent_out,
417            inputs,
418            mode,
419            ctx,
420        );
421    }
422
423    let kind = op
424        .primitive_kind()
425        .ok_or_else(|| missing_primitive_kind(op, ADRuleKind::Transpose))?;
426    let rule = registry::primitive_ad_rule(kind)
427        .ok_or_else(|| registry::missing_rule(kind, ADRuleKind::Transpose))?;
428    let mask = rule.residual_mask();
429    let transpose_inputs = inputs
430        .iter()
431        .enumerate()
432        .map(|(index, input)| transpose_input::TransposeInputRef::new(input, index, mask))
433        .collect::<Vec<_>>();
434    let builder_dyn: &mut dyn PrimitiveRuleBuilder = builder;
435    rule.transpose_rule(op, builder_dyn, cotangent_out, &transpose_inputs, mode, ctx)
436}
437
438#[cfg(test)]
439mod tests;
440
441#[cfg(test)]
442mod residual_spec_tests {
443    use super::ResidualSpec;
444
445    #[test]
446    fn residual_spec_builders_and_queries_cover_always_on_api() {
447        let none = ResidualSpec::none();
448        assert!(none.is_empty());
449        assert!(!none.declares_input(0));
450        assert!(!none.declares_output(0));
451
452        let single = ResidualSpec::input(1).with_output(0);
453        assert!(!single.declares_input(0));
454        assert!(single.declares_input(1));
455        assert!(single.declares_output(0));
456        assert!(!single.is_empty());
457
458        let chained = ResidualSpec::output(2).with_input(0).with_input(3);
459        assert!(chained.declares_input(0) && chained.declares_input(3));
460        assert!(!chained.declares_input(1));
461        assert!(chained.declares_output(2));
462
463        let all_in = ResidualSpec::all_inputs();
464        assert!(all_in.declares_input(0) && all_in.declares_input(63));
465        assert!(!all_in.declares_output(0));
466        let all_out = ResidualSpec::all_outputs();
467        assert!(all_out.declares_output(0) && all_out.declares_output(63));
468        assert!(!all_out.declares_input(0));
469
470        let both = ResidualSpec::all_inputs().with_all_outputs();
471        assert!(both.declares_input(63) && both.declares_output(63));
472        assert!(!both.declares_input(64) && !both.declares_output(64));
473
474        assert!(ResidualSpec::default().is_empty());
475    }
476}
tenferro_ops/ad/mod.rs

tenferro_ops/ad/
mod.rs
