Skip to main content

tenferro_ad/
semantic_extension.rs

1//! Semantic-program automatic-differentiation rules for extension operations.
2//!
3//! Rules in this module are owned explicitly by [`crate::AdContext`]. They
4//! operate on opaque semantic [`ProgramValue`] tokens and never expose
5//! computegraph node keys or execution-program slots.
6
7use std::collections::HashMap;
8use std::fmt::Debug;
9use std::sync::Arc;
10
11pub use tenferro_ops::ad::ResidualSpec;
12use tenferro_ops::ext_op::ExtensionOp;
13use tenferro_runtime::program::{
14    ProgramBuildError, ProgramValue, ProgramValueMetadata, SemanticOpRef, SemanticOperationView,
15    SemanticProgramBuilder, SemanticProvenanceView,
16};
17
18/// One optional semantic AD value.
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub enum AdValue {
21    /// This primal, tangent, or cotangent is inactive.
22    Absent,
23    /// Active value owned by the destination semantic-program builder.
24    Value(ProgramValue),
25}
26
27impl AdValue {
28    /// Return the active semantic value, if present.
29    #[must_use]
30    pub const fn value(self) -> Option<ProgramValue> {
31        match self {
32            Self::Absent => None,
33            Self::Value(value) => Some(value),
34        }
35    }
36}
37
38/// Semantic extension AD rule role.
39#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
40pub enum SemanticAdRuleRole {
41    /// Definitional forward linearization.
42    Linearize,
43    /// Transpose of a linearized extension operation.
44    LinearTranspose,
45    /// Direct reverse rule expressed against primal values.
46    PrimalVjp,
47}
48
49/// Registration failures for semantic extension AD rules.
50#[derive(Debug, thiserror::Error)]
51pub enum SemanticExtensionRegistryError {
52    /// A rule with the same family and role is already present.
53    #[error("semantic extension AD {role:?} rule for family {family_id:?} is already registered")]
54    DuplicateRule {
55        /// Duplicate extension family identifier.
56        family_id: &'static str,
57        /// Duplicate semantic AD role.
58        role: SemanticAdRuleRole,
59    },
60    /// A rule family is not a namespaced, versioned identifier.
61    #[error("semantic extension AD family {family_id:?} is not namespaced and versioned")]
62    MalformedFamilyId {
63        /// Invalid extension family identifier.
64        family_id: &'static str,
65    },
66}
67
68/// Whether a checked semantic AD request refers to a primal input or output.
69///
70/// # Examples
71///
72/// ```
73/// use tenferro_ad::semantic_extension::PrimalValueKind;
74///
75/// assert_eq!(format!("{:?}", PrimalValueKind::Input), "Input");
76/// ```
77#[derive(Clone, Copy, Debug, PartialEq, Eq)]
78#[non_exhaustive]
79pub enum PrimalValueKind {
80    /// An ordered primal input.
81    Input,
82    /// An ordered primal output.
83    Output,
84}
85
86/// Failures while dispatching or building semantic extension AD.
87#[derive(Debug, thiserror::Error)]
88pub enum SemanticAdError {
89    /// The supplied operation is not an extension operation.
90    #[error("semantic AD extension dispatch received a core operation")]
91    CoreOperation,
92    /// Observable effects make implicit differentiation unsafe.
93    #[error("semantic extension family {family_id:?} has observable effects")]
94    EffectfulExtension {
95        /// Effectful extension family.
96        family_id: &'static str,
97    },
98    /// No rule is registered for the requested family and role.
99    #[error("semantic extension family {family_id:?} has no {role:?} AD rule")]
100    MissingRule {
101        /// Extension family without a rule.
102        family_id: &'static str,
103        /// Missing semantic AD role.
104        role: SemanticAdRuleRole,
105    },
106    /// An ordered request or result field has the wrong length.
107    #[error("semantic AD field {field} expects {expected} values, got {actual}")]
108    Arity {
109        /// Name of the invalid ordered field.
110        field: &'static str,
111        /// Required field length.
112        expected: usize,
113        /// Supplied field length.
114        actual: usize,
115    },
116    /// A checked primal value index is outside the request's ordered values.
117    #[error(
118        "semantic extension family {family_id:?} {kind:?} index {index} is out of bounds for length {len}"
119    )]
120    PrimalIndexOutOfBounds {
121        /// Extension family that owns the request.
122        family_id: &'static str,
123        /// Ordered primal collection being accessed.
124        kind: PrimalValueKind,
125        /// Requested index.
126        index: usize,
127        /// Number of values in the collection.
128        len: usize,
129    },
130    /// A checked primal value was not declared as a tensor residual by the rule.
131    #[error(
132        "semantic extension family {family_id:?} may not access undeclared {kind:?} residual value {index}"
133    )]
134    UndeclaredResidualValue {
135        /// Extension family that owns the request.
136        family_id: &'static str,
137        /// Ordered primal collection being accessed.
138        kind: PrimalValueKind,
139        /// Requested index.
140        index: usize,
141    },
142    /// A request or result value belongs to another builder.
143    #[error("semantic AD field {field}[{index}] does not belong to the destination builder")]
144    ForeignValue {
145        /// Name of the invalid ordered field.
146        field: &'static str,
147        /// Index of the foreign value.
148        index: usize,
149    },
150    /// A family-specific rule deliberately rejects this payload.
151    #[error("semantic extension family {family_id:?} does not support {role:?}: {message}")]
152    Unsupported {
153        /// Extension family that rejected the transform.
154        family_id: &'static str,
155        /// Rejected semantic AD role.
156        role: SemanticAdRuleRole,
157        /// Bounded family-specific diagnostic.
158        message: String,
159    },
160    /// A family-specific rule failed with a typed source error.
161    #[error("semantic extension family {family_id:?} {role:?} rule failed: {source}")]
162    Rule {
163        /// Extension family whose rule failed.
164        family_id: &'static str,
165        /// Semantic AD role being evaluated.
166        role: SemanticAdRuleRole,
167        /// Original typed rule failure.
168        #[source]
169        source: Box<dyn std::error::Error + Send + Sync + 'static>,
170    },
171    /// A family-specific semantic rule invariant was violated.
172    #[error("semantic extension family {family_id:?} {role:?} invariant failed: {message}")]
173    Invariant {
174        /// Extension family whose invariant failed.
175        family_id: &'static str,
176        /// Semantic AD role being evaluated.
177        role: SemanticAdRuleRole,
178        /// Bounded invariant diagnostic.
179        message: String,
180    },
181    /// Semantic-program construction failed inside a rule.
182    #[error("semantic extension AD program construction failed: {0}")]
183    Build(#[from] ProgramBuildError),
184}
185
186/// Ordered inputs for one semantic extension linearization rule.
187#[derive(Clone, Copy)]
188pub struct SemanticLinearizeRequest<'a> {
189    op: &'a dyn ExtensionOp,
190    primal_inputs: &'a [ProgramValue],
191    primal_outputs: &'a [ProgramValue],
192    tangent_inputs: &'a [AdValue],
193    active_outputs: &'a [bool],
194    provenance: SemanticProvenanceView<'a>,
195}
196
197impl<'a> SemanticLinearizeRequest<'a> {
198    /// Borrow the extension payload.
199    pub const fn op(self) -> &'a dyn ExtensionOp {
200        self.op
201    }
202
203    /// Borrow ordered destination-local primal inputs.
204    pub const fn primal_inputs(self) -> &'a [ProgramValue] {
205        self.primal_inputs
206    }
207
208    /// Borrow ordered destination-local primal outputs.
209    pub const fn primal_outputs(self) -> &'a [ProgramValue] {
210        self.primal_outputs
211    }
212
213    /// Borrow ordered optional tangent inputs.
214    pub const fn tangent_inputs(self) -> &'a [AdValue] {
215        self.tangent_inputs
216    }
217
218    /// Borrow the ordered active-output mask.
219    pub const fn active_outputs(self) -> &'a [bool] {
220        self.active_outputs
221    }
222
223    /// Return bounded operation provenance.
224    pub const fn provenance(self) -> SemanticProvenanceView<'a> {
225        self.provenance
226    }
227}
228
229/// Output of one semantic extension linearization rule.
230#[derive(Clone, Debug, PartialEq, Eq)]
231pub struct SemanticLinearizeResult {
232    tangent_outputs: Box<[AdValue]>,
233    residuals: Box<[ProgramValue]>,
234}
235
236impl SemanticLinearizeResult {
237    /// Construct ordered tangent outputs and residuals.
238    #[must_use]
239    pub fn new(
240        tangent_outputs: impl IntoIterator<Item = AdValue>,
241        residuals: impl IntoIterator<Item = ProgramValue>,
242    ) -> Self {
243        Self {
244            tangent_outputs: tangent_outputs.into_iter().collect(),
245            residuals: residuals.into_iter().collect(),
246        }
247    }
248
249    /// Borrow ordered optional tangent outputs.
250    pub fn tangent_outputs(&self) -> &[AdValue] {
251        &self.tangent_outputs
252    }
253
254    /// Borrow ordered residual values saved for transpose.
255    pub fn residuals(&self) -> &[ProgramValue] {
256        &self.residuals
257    }
258}
259
260/// Ordered inputs for one semantic linear-transpose rule.
261pub struct SemanticLinearTransposeRequest<'a> {
262    op: &'a dyn ExtensionOp,
263    primal_inputs: &'a [ProgramValue],
264    primal_outputs: &'a [ProgramValue],
265    primal_input_metadata: Box<[ProgramValueMetadata]>,
266    primal_output_metadata: Box<[ProgramValueMetadata]>,
267    cotangent_outputs: &'a [AdValue],
268    active_inputs: &'a [bool],
269    residuals: &'a [ProgramValue],
270    residual_mask: ResidualSpec,
271    provenance: SemanticProvenanceView<'a>,
272}
273
274impl<'a> SemanticLinearTransposeRequest<'a> {
275    /// Borrow the extension payload.
276    ///
277    /// # Examples
278    ///
279    /// ```rust
280    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
281    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
282    /// let _ = request.op().family_id();
283    /// # Ok(())
284    /// # }
285    /// ```
286    pub fn op(&self) -> &dyn ExtensionOp {
287        self.op
288    }
289
290    /// Return one declared primal input tensor value.
291    ///
292    /// # Errors
293    ///
294    /// Returns [`SemanticAdError::PrimalIndexOutOfBounds`] before checking the
295    /// residual mask, or [`SemanticAdError::UndeclaredResidualValue`] when the
296    /// rule did not declare this input as a tensor residual.
297    ///
298    /// # Examples
299    ///
300    /// ```rust
301    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
302    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
303    /// let _ = request.primal_input_value(0)?;
304    /// # Ok(())
305    /// # }
306    /// ```
307    pub fn primal_input_value(&self, index: usize) -> Result<ProgramValue, SemanticAdError> {
308        checked_primal_value(
309            self.op.family_id(),
310            PrimalValueKind::Input,
311            index,
312            self.primal_inputs,
313            self.residual_mask,
314        )
315    }
316
317    /// Return one declared primal output tensor value.
318    ///
319    /// # Errors
320    ///
321    /// Returns [`SemanticAdError::PrimalIndexOutOfBounds`] before checking the
322    /// residual mask, or [`SemanticAdError::UndeclaredResidualValue`] when the
323    /// rule did not declare this output as a tensor residual.
324    ///
325    /// # Examples
326    ///
327    /// ```rust
328    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
329    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
330    /// let _ = request.primal_output_value(0)?;
331    /// # Ok(())
332    /// # }
333    /// ```
334    pub fn primal_output_value(&self, index: usize) -> Result<ProgramValue, SemanticAdError> {
335        checked_primal_value(
336            self.op.family_id(),
337            PrimalValueKind::Output,
338            index,
339            self.primal_outputs,
340            self.residual_mask,
341        )
342    }
343
344    /// Borrow metadata for one primal input without exposing its value token.
345    ///
346    /// # Errors
347    ///
348    /// Returns [`SemanticAdError::PrimalIndexOutOfBounds`] when `index` is not
349    /// an ordered primal input.
350    ///
351    /// # Examples
352    ///
353    /// ```rust
354    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
355    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
356    /// let _ = request.primal_input_meta(0)?;
357    /// # Ok(())
358    /// # }
359    /// ```
360    pub fn primal_input_meta(
361        &self,
362        index: usize,
363    ) -> Result<&ProgramValueMetadata, SemanticAdError> {
364        checked_primal_metadata(
365            self.op.family_id(),
366            PrimalValueKind::Input,
367            index,
368            &self.primal_input_metadata,
369        )
370    }
371
372    /// Borrow metadata for one primal output without exposing its value token.
373    ///
374    /// # Errors
375    ///
376    /// Returns [`SemanticAdError::PrimalIndexOutOfBounds`] when `index` is not
377    /// an ordered primal output.
378    ///
379    /// # Examples
380    ///
381    /// ```rust
382    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
383    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
384    /// let _ = request.primal_output_meta(0)?;
385    /// # Ok(())
386    /// # }
387    /// ```
388    pub fn primal_output_meta(
389        &self,
390        index: usize,
391    ) -> Result<&ProgramValueMetadata, SemanticAdError> {
392        checked_primal_metadata(
393            self.op.family_id(),
394            PrimalValueKind::Output,
395            index,
396            &self.primal_output_metadata,
397        )
398    }
399
400    /// Return the number of ordered primal inputs.
401    #[must_use]
402    ///
403    /// # Examples
404    ///
405    /// ```rust
406    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
407    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
408    /// let _ = request.primal_input_count();
409    /// # Ok(())
410    /// # }
411    /// ```
412    pub const fn primal_input_count(&self) -> usize {
413        self.primal_input_metadata.len()
414    }
415
416    /// Return the number of ordered primal outputs.
417    #[must_use]
418    ///
419    /// # Examples
420    ///
421    /// ```rust
422    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
423    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
424    /// let _ = request.primal_output_count();
425    /// # Ok(())
426    /// # }
427    /// ```
428    pub const fn primal_output_count(&self) -> usize {
429        self.primal_output_metadata.len()
430    }
431
432    /// Borrow ordered optional output cotangents.
433    ///
434    /// # Examples
435    ///
436    /// ```rust
437    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
438    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
439    /// let _ = request.cotangent_outputs();
440    /// # Ok(())
441    /// # }
442    /// ```
443    pub fn cotangent_outputs(&self) -> &[AdValue] {
444        self.cotangent_outputs
445    }
446
447    /// Borrow the ordered active-input mask.
448    ///
449    /// # Examples
450    ///
451    /// ```rust
452    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
453    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
454    /// let _ = request.active_inputs();
455    /// # Ok(())
456    /// # }
457    /// ```
458    pub fn active_inputs(&self) -> &[bool] {
459        self.active_inputs
460    }
461
462    /// Borrow ordered residuals produced by linearization.
463    ///
464    /// # Examples
465    ///
466    /// ```rust
467    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
468    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
469    /// let _ = request.residuals();
470    /// # Ok(())
471    /// # }
472    /// ```
473    pub fn residuals(&self) -> &[ProgramValue] {
474        self.residuals
475    }
476
477    /// Return this rule's declared residual mask: which primal input/output
478    /// indices may be read as tensor values. Accesses outside the mask must
479    /// only use shape/dtype metadata.
480    ///
481    /// # Examples
482    ///
483    /// ```rust
484    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
485    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
486    /// let _ = request.residual_mask();
487    /// # Ok(())
488    /// # }
489    /// ```
490    pub const fn residual_mask(&self) -> ResidualSpec {
491        self.residual_mask
492    }
493
494    /// Return bounded operation provenance.
495    ///
496    /// # Examples
497    ///
498    /// ```rust
499    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticLinearTransposeRequest};
500    /// # fn inspect(request: &SemanticLinearTransposeRequest<'_>) -> Result<(), SemanticAdError> {
501    /// let _ = request.provenance();
502    /// # Ok(())
503    /// # }
504    /// ```
505    pub const fn provenance(&self) -> SemanticProvenanceView<'a> {
506        self.provenance
507    }
508}
509
510/// Ordered inputs for one direct semantic primal-VJP rule.
511pub struct SemanticPrimalVjpRequest<'a> {
512    op: &'a dyn ExtensionOp,
513    primal_inputs: &'a [ProgramValue],
514    primal_outputs: &'a [ProgramValue],
515    primal_input_metadata: Box<[ProgramValueMetadata]>,
516    primal_output_metadata: Box<[ProgramValueMetadata]>,
517    cotangent_outputs: &'a [AdValue],
518    active_inputs: &'a [bool],
519    residual_mask: ResidualSpec,
520    provenance: SemanticProvenanceView<'a>,
521}
522
523impl<'a> SemanticPrimalVjpRequest<'a> {
524    /// Borrow the extension payload.
525    ///
526    /// # Examples
527    ///
528    /// ```rust
529    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
530    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
531    /// let _ = request.op().family_id();
532    /// # Ok(())
533    /// # }
534    /// ```
535    pub fn op(&self) -> &dyn ExtensionOp {
536        self.op
537    }
538
539    /// Return one declared primal input tensor value.
540    ///
541    /// # Errors
542    ///
543    /// Returns [`SemanticAdError::PrimalIndexOutOfBounds`] before checking the
544    /// residual mask, or [`SemanticAdError::UndeclaredResidualValue`] when the
545    /// rule did not declare this input as a tensor residual.
546    ///
547    /// # Examples
548    ///
549    /// ```rust
550    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
551    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
552    /// let _ = request.primal_input_value(0)?;
553    /// # Ok(())
554    /// # }
555    /// ```
556    pub fn primal_input_value(&self, index: usize) -> Result<ProgramValue, SemanticAdError> {
557        checked_primal_value(
558            self.op.family_id(),
559            PrimalValueKind::Input,
560            index,
561            self.primal_inputs,
562            self.residual_mask,
563        )
564    }
565
566    /// Return one declared primal output tensor value.
567    ///
568    /// # Errors
569    ///
570    /// Returns [`SemanticAdError::PrimalIndexOutOfBounds`] before checking the
571    /// residual mask, or [`SemanticAdError::UndeclaredResidualValue`] when the
572    /// rule did not declare this output as a tensor residual.
573    ///
574    /// # Examples
575    ///
576    /// ```rust
577    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
578    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
579    /// let _ = request.primal_output_value(0)?;
580    /// # Ok(())
581    /// # }
582    /// ```
583    pub fn primal_output_value(&self, index: usize) -> Result<ProgramValue, SemanticAdError> {
584        checked_primal_value(
585            self.op.family_id(),
586            PrimalValueKind::Output,
587            index,
588            self.primal_outputs,
589            self.residual_mask,
590        )
591    }
592
593    /// Borrow metadata for one primal input without exposing its value token.
594    ///
595    /// # Errors
596    ///
597    /// Returns [`SemanticAdError::PrimalIndexOutOfBounds`] when `index` is not
598    /// an ordered primal input.
599    ///
600    /// # Examples
601    ///
602    /// ```rust
603    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
604    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
605    /// let _ = request.primal_input_meta(0)?;
606    /// # Ok(())
607    /// # }
608    /// ```
609    pub fn primal_input_meta(
610        &self,
611        index: usize,
612    ) -> Result<&ProgramValueMetadata, SemanticAdError> {
613        checked_primal_metadata(
614            self.op.family_id(),
615            PrimalValueKind::Input,
616            index,
617            &self.primal_input_metadata,
618        )
619    }
620
621    /// Borrow metadata for one primal output without exposing its value token.
622    ///
623    /// # Errors
624    ///
625    /// Returns [`SemanticAdError::PrimalIndexOutOfBounds`] when `index` is not
626    /// an ordered primal output.
627    ///
628    /// # Examples
629    ///
630    /// ```rust
631    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
632    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
633    /// let _ = request.primal_output_meta(0)?;
634    /// # Ok(())
635    /// # }
636    /// ```
637    pub fn primal_output_meta(
638        &self,
639        index: usize,
640    ) -> Result<&ProgramValueMetadata, SemanticAdError> {
641        checked_primal_metadata(
642            self.op.family_id(),
643            PrimalValueKind::Output,
644            index,
645            &self.primal_output_metadata,
646        )
647    }
648
649    /// Return the number of ordered primal inputs.
650    #[must_use]
651    ///
652    /// # Examples
653    ///
654    /// ```rust
655    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
656    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
657    /// let _ = request.primal_input_count();
658    /// # Ok(())
659    /// # }
660    /// ```
661    pub const fn primal_input_count(&self) -> usize {
662        self.primal_input_metadata.len()
663    }
664
665    /// Return the number of ordered primal outputs.
666    #[must_use]
667    ///
668    /// # Examples
669    ///
670    /// ```rust
671    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
672    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
673    /// let _ = request.primal_output_count();
674    /// # Ok(())
675    /// # }
676    /// ```
677    pub const fn primal_output_count(&self) -> usize {
678        self.primal_output_metadata.len()
679    }
680
681    /// Borrow ordered optional output cotangents.
682    ///
683    /// # Examples
684    ///
685    /// ```rust
686    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
687    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
688    /// let _ = request.cotangent_outputs();
689    /// # Ok(())
690    /// # }
691    /// ```
692    pub fn cotangent_outputs(&self) -> &[AdValue] {
693        self.cotangent_outputs
694    }
695
696    /// Borrow the ordered active-input mask.
697    ///
698    /// # Examples
699    ///
700    /// ```rust
701    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
702    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
703    /// let _ = request.active_inputs();
704    /// # Ok(())
705    /// # }
706    /// ```
707    pub fn active_inputs(&self) -> &[bool] {
708        self.active_inputs
709    }
710
711    /// Return this rule's declared residual mask: which primal input/output
712    /// indices may be read as tensor values. Accesses outside the mask must
713    /// only use shape/dtype metadata.
714    ///
715    /// # Examples
716    ///
717    /// ```rust
718    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
719    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
720    /// let _ = request.residual_mask();
721    /// # Ok(())
722    /// # }
723    /// ```
724    pub const fn residual_mask(&self) -> ResidualSpec {
725        self.residual_mask
726    }
727
728    /// Return bounded operation provenance.
729    ///
730    /// # Examples
731    ///
732    /// ```rust
733    /// # use tenferro_ad::semantic_extension::{SemanticAdError, SemanticPrimalVjpRequest};
734    /// # fn inspect(request: &SemanticPrimalVjpRequest<'_>) -> Result<(), SemanticAdError> {
735    /// let _ = request.provenance();
736    /// # Ok(())
737    /// # }
738    /// ```
739    pub const fn provenance(&self) -> SemanticProvenanceView<'a> {
740        self.provenance
741    }
742}
743
744/// Definitional JVP rule for one extension family.
745pub trait SemanticLinearizeRule: Debug + Send + Sync + 'static {
746    /// Return the versioned extension family handled by this rule.
747    fn family_id(&self) -> &'static str;
748
749    /// Emit ordered tangent outputs and residuals into `builder`.
750    ///
751    /// # Errors
752    ///
753    /// Returns [`SemanticAdError::Unsupported`] when the payload is outside
754    /// the rule's supported domain, or [`SemanticAdError::Build`] when emitted
755    /// semantic operations fail validation.
756    fn linearize(
757        &self,
758        request: SemanticLinearizeRequest<'_>,
759        builder: &mut SemanticProgramBuilder,
760    ) -> Result<SemanticLinearizeResult, SemanticAdError>;
761}
762
763/// Transpose rule for an extension viewed as a linear map.
764pub trait SemanticLinearTransposeRule: Debug + Send + Sync + 'static {
765    /// Return the versioned extension family handled by this rule.
766    fn family_id(&self) -> &'static str;
767
768    /// Declare which primal input/output indices this rule reads as tensor
769    /// residuals. Indices not declared may only be accessed through metadata.
770    fn residual_mask(&self) -> ResidualSpec;
771
772    /// Emit ordered optional input cotangents into `builder`.
773    ///
774    /// # Errors
775    ///
776    /// Returns [`SemanticAdError::Unsupported`] when the payload is outside
777    /// the rule's supported domain, or [`SemanticAdError::Build`] when emitted
778    /// semantic operations fail validation.
779    fn linear_transpose(
780        &self,
781        request: SemanticLinearTransposeRequest<'_>,
782        builder: &mut SemanticProgramBuilder,
783    ) -> Result<Box<[AdValue]>, SemanticAdError>;
784}
785
786/// Optional direct VJP rule expressed against primal semantic values.
787pub trait SemanticPrimalVjpRule: Debug + Send + Sync + 'static {
788    /// Return the versioned extension family handled by this rule.
789    fn family_id(&self) -> &'static str;
790
791    /// Declare which primal input/output indices this rule reads as tensor
792    /// residuals. Indices not declared may only be accessed through metadata.
793    fn residual_mask(&self) -> ResidualSpec;
794
795    /// Emit ordered optional input cotangents into `builder`.
796    ///
797    /// # Errors
798    ///
799    /// Returns [`SemanticAdError::Unsupported`] when the payload is outside
800    /// the rule's supported domain, or [`SemanticAdError::Build`] when emitted
801    /// semantic operations fail validation.
802    fn primal_vjp(
803        &self,
804        request: SemanticPrimalVjpRequest<'_>,
805        builder: &mut SemanticProgramBuilder,
806    ) -> Result<Box<[AdValue]>, SemanticAdError>;
807}
808
809type LinearizeMap = HashMap<&'static str, Arc<dyn SemanticLinearizeRule>>;
810type LinearTransposeMap = HashMap<&'static str, Arc<dyn SemanticLinearTransposeRule>>;
811type PrimalVjpMap = HashMap<&'static str, Arc<dyn SemanticPrimalVjpRule>>;
812
813/// Explicit clone-on-write set of semantic extension AD rules.
814#[derive(Clone, Default)]
815pub struct SemanticExtensionRuleSet {
816    linearize: Arc<LinearizeMap>,
817    linear_transpose: Arc<LinearTransposeMap>,
818    primal_vjp: Arc<PrimalVjpMap>,
819}
820
821impl Debug for SemanticExtensionRuleSet {
822    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
823        let mut linearize: Vec<_> = self.linearize.keys().copied().collect();
824        let mut linear_transpose: Vec<_> = self.linear_transpose.keys().copied().collect();
825        let mut primal_vjp: Vec<_> = self.primal_vjp.keys().copied().collect();
826        linearize.sort_unstable();
827        linear_transpose.sort_unstable();
828        primal_vjp.sort_unstable();
829        formatter
830            .debug_struct("SemanticExtensionRuleSet")
831            .field("linearize", &linearize)
832            .field("linear_transpose", &linear_transpose)
833            .field("primal_vjp", &primal_vjp)
834            .finish()
835    }
836}
837
838impl SemanticExtensionRuleSet {
839    /// Construct an empty rule set.
840    #[must_use]
841    pub fn new() -> Self {
842        Self::default()
843    }
844
845    /// Register one semantic linearize rule.
846    ///
847    /// # Errors
848    ///
849    /// Returns [`SemanticExtensionRegistryError::MalformedFamilyId`] for an
850    /// invalid family or [`SemanticExtensionRegistryError::DuplicateRule`] for
851    /// an existing family in this role.
852    pub fn register_linearize(
853        &mut self,
854        rule: Arc<dyn SemanticLinearizeRule>,
855    ) -> Result<(), SemanticExtensionRegistryError> {
856        validate_insert(
857            &self.linearize,
858            rule.family_id(),
859            SemanticAdRuleRole::Linearize,
860        )?;
861        Arc::make_mut(&mut self.linearize).insert(rule.family_id(), rule);
862        Ok(())
863    }
864
865    /// Register one semantic linear-transpose rule.
866    ///
867    /// # Errors
868    ///
869    /// Returns [`SemanticExtensionRegistryError::MalformedFamilyId`] for an
870    /// invalid family or [`SemanticExtensionRegistryError::DuplicateRule`] for
871    /// an existing family in this role.
872    pub fn register_linear_transpose(
873        &mut self,
874        rule: Arc<dyn SemanticLinearTransposeRule>,
875    ) -> Result<(), SemanticExtensionRegistryError> {
876        validate_insert(
877            &self.linear_transpose,
878            rule.family_id(),
879            SemanticAdRuleRole::LinearTranspose,
880        )?;
881        Arc::make_mut(&mut self.linear_transpose).insert(rule.family_id(), rule);
882        Ok(())
883    }
884
885    /// Register one direct semantic primal-VJP rule.
886    ///
887    /// # Errors
888    ///
889    /// Returns [`SemanticExtensionRegistryError::MalformedFamilyId`] for an
890    /// invalid family or [`SemanticExtensionRegistryError::DuplicateRule`] for
891    /// an existing family in this role.
892    pub fn register_primal_vjp(
893        &mut self,
894        rule: Arc<dyn SemanticPrimalVjpRule>,
895    ) -> Result<(), SemanticExtensionRegistryError> {
896        validate_insert(
897            &self.primal_vjp,
898            rule.family_id(),
899            SemanticAdRuleRole::PrimalVjp,
900        )?;
901        Arc::make_mut(&mut self.primal_vjp).insert(rule.family_id(), rule);
902        Ok(())
903    }
904
905    /// Return a rule set containing one semantic linearize rule.
906    ///
907    /// # Errors
908    ///
909    /// Returns [`SemanticExtensionRegistryError::MalformedFamilyId`] for an
910    /// invalid family or [`SemanticExtensionRegistryError::DuplicateRule`] for
911    /// an existing linearize rule.
912    pub fn with_linearize(
913        mut self,
914        rule: Arc<dyn SemanticLinearizeRule>,
915    ) -> Result<Self, SemanticExtensionRegistryError> {
916        self.register_linearize(rule)?;
917        Ok(self)
918    }
919
920    /// Return a rule set containing one semantic linear-transpose rule.
921    ///
922    /// # Errors
923    ///
924    /// Returns [`SemanticExtensionRegistryError::MalformedFamilyId`] for an
925    /// invalid family or [`SemanticExtensionRegistryError::DuplicateRule`] for
926    /// an existing linear-transpose rule.
927    pub fn with_linear_transpose(
928        mut self,
929        rule: Arc<dyn SemanticLinearTransposeRule>,
930    ) -> Result<Self, SemanticExtensionRegistryError> {
931        self.register_linear_transpose(rule)?;
932        Ok(self)
933    }
934
935    /// Return a rule set containing one direct semantic primal-VJP rule.
936    ///
937    /// # Errors
938    ///
939    /// Returns [`SemanticExtensionRegistryError::MalformedFamilyId`] for an
940    /// invalid family or [`SemanticExtensionRegistryError::DuplicateRule`] for
941    /// an existing primal-VJP rule.
942    pub fn with_primal_vjp(
943        mut self,
944        rule: Arc<dyn SemanticPrimalVjpRule>,
945    ) -> Result<Self, SemanticExtensionRegistryError> {
946        self.register_primal_vjp(rule)?;
947        Ok(self)
948    }
949
950    /// Merge another rule set atomically.
951    ///
952    /// # Errors
953    ///
954    /// Returns [`SemanticExtensionRegistryError::MalformedFamilyId`] for an
955    /// invalid family or [`SemanticExtensionRegistryError::DuplicateRule`] for
956    /// a role-equivalent duplicate. The receiver is unchanged on failure.
957    pub fn merge(&mut self, other: Self) -> Result<(), SemanticExtensionRegistryError> {
958        let mut candidate = self.clone();
959        for rule in other.linearize.values() {
960            candidate.register_linearize(Arc::clone(rule))?;
961        }
962        for rule in other.linear_transpose.values() {
963            candidate.register_linear_transpose(Arc::clone(rule))?;
964        }
965        for rule in other.primal_vjp.values() {
966            candidate.register_primal_vjp(Arc::clone(rule))?;
967        }
968        *self = candidate;
969        Ok(())
970    }
971
972    /// Look up a semantic linearize rule by extension family.
973    #[must_use]
974    pub fn lookup_linearize(&self, family_id: &str) -> Option<Arc<dyn SemanticLinearizeRule>> {
975        self.linearize.get(family_id).cloned()
976    }
977
978    /// Look up a semantic linear-transpose rule by extension family.
979    #[must_use]
980    pub fn lookup_linear_transpose(
981        &self,
982        family_id: &str,
983    ) -> Option<Arc<dyn SemanticLinearTransposeRule>> {
984        self.linear_transpose.get(family_id).cloned()
985    }
986
987    /// Look up a direct semantic primal-VJP rule by extension family.
988    #[must_use]
989    pub fn lookup_primal_vjp(&self, family_id: &str) -> Option<Arc<dyn SemanticPrimalVjpRule>> {
990        self.primal_vjp.get(family_id).cloned()
991    }
992
993    /// Validate and dispatch one semantic extension linearization.
994    ///
995    /// # Errors
996    ///
997    /// Returns [`SemanticAdError::CoreOperation`] for a core operation,
998    /// [`SemanticAdError::EffectfulExtension`] before rule dispatch for an
999    /// effectful extension, or typed rule/arity/ownership/build failures.
1000    #[allow(clippy::too_many_arguments)]
1001    pub fn linearize_operation(
1002        &self,
1003        operation: SemanticOperationView<'_>,
1004        primal_inputs: &[ProgramValue],
1005        primal_outputs: &[ProgramValue],
1006        tangent_inputs: &[AdValue],
1007        active_outputs: &[bool],
1008        builder: &mut SemanticProgramBuilder,
1009    ) -> Result<SemanticLinearizeResult, SemanticAdError> {
1010        let op = extension_for_dispatch(operation)?;
1011        validate_operation_inputs(operation, primal_inputs, primal_outputs, builder)?;
1012        validate_len("tangent_inputs", op.input_count(), tangent_inputs.len())?;
1013        validate_len("active_outputs", op.output_count(), active_outputs.len())?;
1014        validate_ad_values("tangent_inputs", tangent_inputs, builder)?;
1015        let rule = self
1016            .lookup_linearize(op.family_id())
1017            .ok_or(SemanticAdError::MissingRule {
1018                family_id: op.family_id(),
1019                role: SemanticAdRuleRole::Linearize,
1020            })?;
1021        let result = rule.linearize(
1022            SemanticLinearizeRequest {
1023                op,
1024                primal_inputs,
1025                primal_outputs,
1026                tangent_inputs,
1027                active_outputs,
1028                provenance: operation.provenance(),
1029            },
1030            builder,
1031        )?;
1032        validate_len(
1033            "tangent_outputs",
1034            op.output_count(),
1035            result.tangent_outputs.len(),
1036        )?;
1037        validate_ad_values("tangent_outputs", &result.tangent_outputs, builder)?;
1038        validate_values("residuals", &result.residuals, builder)?;
1039        Ok(result)
1040    }
1041
1042    /// Validate and dispatch one semantic extension linear transpose.
1043    ///
1044    /// # Errors
1045    ///
1046    /// Returns [`SemanticAdError::EffectfulExtension`] before rule dispatch,
1047    /// [`SemanticAdError::MissingRule`] when no transpose rule exists,
1048    /// [`SemanticAdError::Arity`] / [`SemanticAdError::ForeignValue`] for an
1049    /// invalid request or result, or a typed family-rule failure.
1050    #[allow(clippy::too_many_arguments)]
1051    pub fn linear_transpose_operation(
1052        &self,
1053        operation: SemanticOperationView<'_>,
1054        primal_inputs: &[ProgramValue],
1055        primal_outputs: &[ProgramValue],
1056        cotangent_outputs: &[AdValue],
1057        active_inputs: &[bool],
1058        residuals: &[ProgramValue],
1059        builder: &mut SemanticProgramBuilder,
1060    ) -> Result<Box<[AdValue]>, SemanticAdError> {
1061        let op = extension_for_dispatch(operation)?;
1062        validate_operation_inputs(operation, primal_inputs, primal_outputs, builder)?;
1063        validate_len(
1064            "cotangent_outputs",
1065            op.output_count(),
1066            cotangent_outputs.len(),
1067        )?;
1068        validate_len("active_inputs", op.input_count(), active_inputs.len())?;
1069        validate_ad_values("cotangent_outputs", cotangent_outputs, builder)?;
1070        validate_values("residuals", residuals, builder)?;
1071        let primal_input_metadata = snapshot_metadata(primal_inputs, builder)?;
1072        let primal_output_metadata = snapshot_metadata(primal_outputs, builder)?;
1073        let rule =
1074            self.lookup_linear_transpose(op.family_id())
1075                .ok_or(SemanticAdError::MissingRule {
1076                    family_id: op.family_id(),
1077                    role: SemanticAdRuleRole::LinearTranspose,
1078                })?;
1079        let result = rule.linear_transpose(
1080            SemanticLinearTransposeRequest {
1081                op,
1082                primal_inputs,
1083                primal_outputs,
1084                primal_input_metadata,
1085                primal_output_metadata,
1086                cotangent_outputs,
1087                active_inputs,
1088                residuals,
1089                residual_mask: rule.residual_mask(),
1090                provenance: operation.provenance(),
1091            },
1092            builder,
1093        )?;
1094        validate_len("cotangent_inputs", op.input_count(), result.len())?;
1095        validate_ad_values("cotangent_inputs", &result, builder)?;
1096        Ok(result)
1097    }
1098
1099    /// Validate and dispatch one direct semantic primal VJP.
1100    ///
1101    /// # Errors
1102    ///
1103    /// Returns [`SemanticAdError::EffectfulExtension`] before rule dispatch,
1104    /// [`SemanticAdError::MissingRule`] when no primal-VJP rule exists,
1105    /// [`SemanticAdError::Arity`] / [`SemanticAdError::ForeignValue`] for an
1106    /// invalid request or result, or a typed family-rule failure.
1107    #[allow(clippy::too_many_arguments)]
1108    pub fn primal_vjp_operation(
1109        &self,
1110        operation: SemanticOperationView<'_>,
1111        primal_inputs: &[ProgramValue],
1112        primal_outputs: &[ProgramValue],
1113        cotangent_outputs: &[AdValue],
1114        active_inputs: &[bool],
1115        builder: &mut SemanticProgramBuilder,
1116    ) -> Result<Box<[AdValue]>, SemanticAdError> {
1117        let op = extension_for_dispatch(operation)?;
1118        validate_operation_inputs(operation, primal_inputs, primal_outputs, builder)?;
1119        validate_len(
1120            "cotangent_outputs",
1121            op.output_count(),
1122            cotangent_outputs.len(),
1123        )?;
1124        validate_len("active_inputs", op.input_count(), active_inputs.len())?;
1125        validate_ad_values("cotangent_outputs", cotangent_outputs, builder)?;
1126        let primal_input_metadata = snapshot_metadata(primal_inputs, builder)?;
1127        let primal_output_metadata = snapshot_metadata(primal_outputs, builder)?;
1128        let rule = self
1129            .lookup_primal_vjp(op.family_id())
1130            .ok_or(SemanticAdError::MissingRule {
1131                family_id: op.family_id(),
1132                role: SemanticAdRuleRole::PrimalVjp,
1133            })?;
1134        let result = rule.primal_vjp(
1135            SemanticPrimalVjpRequest {
1136                op,
1137                primal_inputs,
1138                primal_outputs,
1139                primal_input_metadata,
1140                primal_output_metadata,
1141                cotangent_outputs,
1142                active_inputs,
1143                residual_mask: rule.residual_mask(),
1144                provenance: operation.provenance(),
1145            },
1146            builder,
1147        )?;
1148        validate_len("cotangent_inputs", op.input_count(), result.len())?;
1149        validate_ad_values("cotangent_inputs", &result, builder)?;
1150        Ok(result)
1151    }
1152}
1153
1154fn extension_for_dispatch(
1155    operation: SemanticOperationView<'_>,
1156) -> Result<&dyn ExtensionOp, SemanticAdError> {
1157    let SemanticOpRef::Extension(op) = operation.op() else {
1158        return Err(SemanticAdError::CoreOperation);
1159    };
1160    if !operation.effects().is_empty() {
1161        return Err(SemanticAdError::EffectfulExtension {
1162            family_id: op.family_id(),
1163        });
1164    }
1165    Ok(op)
1166}
1167
1168fn validate_operation_inputs(
1169    operation: SemanticOperationView<'_>,
1170    primal_inputs: &[ProgramValue],
1171    primal_outputs: &[ProgramValue],
1172    builder: &SemanticProgramBuilder,
1173) -> Result<(), SemanticAdError> {
1174    validate_len(
1175        "primal_inputs",
1176        operation.inputs().len(),
1177        primal_inputs.len(),
1178    )?;
1179    validate_len(
1180        "primal_outputs",
1181        operation.outputs().len(),
1182        primal_outputs.len(),
1183    )?;
1184    validate_values("primal_inputs", primal_inputs, builder)?;
1185    validate_values("primal_outputs", primal_outputs, builder)
1186}
1187
1188fn checked_primal_value(
1189    family_id: &'static str,
1190    kind: PrimalValueKind,
1191    index: usize,
1192    values: &[ProgramValue],
1193    residual_mask: ResidualSpec,
1194) -> Result<ProgramValue, SemanticAdError> {
1195    if index >= values.len() {
1196        return Err(SemanticAdError::PrimalIndexOutOfBounds {
1197            family_id,
1198            kind,
1199            index,
1200            len: values.len(),
1201        });
1202    }
1203    let declared = match kind {
1204        PrimalValueKind::Input => residual_mask.declares_input(index),
1205        PrimalValueKind::Output => residual_mask.declares_output(index),
1206    };
1207    if !declared {
1208        return Err(SemanticAdError::UndeclaredResidualValue {
1209            family_id,
1210            kind,
1211            index,
1212        });
1213    }
1214    Ok(values[index])
1215}
1216
1217fn checked_primal_metadata<'a>(
1218    family_id: &'static str,
1219    kind: PrimalValueKind,
1220    index: usize,
1221    metadata: &'a [ProgramValueMetadata],
1222) -> Result<&'a ProgramValueMetadata, SemanticAdError> {
1223    metadata
1224        .get(index)
1225        .ok_or(SemanticAdError::PrimalIndexOutOfBounds {
1226            family_id,
1227            kind,
1228            index,
1229            len: metadata.len(),
1230        })
1231}
1232
1233fn snapshot_metadata(
1234    values: &[ProgramValue],
1235    builder: &SemanticProgramBuilder,
1236) -> Result<Box<[ProgramValueMetadata]>, SemanticAdError> {
1237    values
1238        .iter()
1239        .copied()
1240        .map(|value| Ok(builder.value_metadata(value)?.clone()))
1241        .collect()
1242}
1243
1244fn validate_values(
1245    field: &'static str,
1246    values: &[ProgramValue],
1247    builder: &SemanticProgramBuilder,
1248) -> Result<(), SemanticAdError> {
1249    for (index, value) in values.iter().copied().enumerate() {
1250        if builder.validate_value(value).is_err() {
1251            return Err(SemanticAdError::ForeignValue { field, index });
1252        }
1253    }
1254    Ok(())
1255}
1256
1257fn validate_ad_values(
1258    field: &'static str,
1259    values: &[AdValue],
1260    builder: &SemanticProgramBuilder,
1261) -> Result<(), SemanticAdError> {
1262    for (index, value) in values.iter().copied().enumerate() {
1263        if let AdValue::Value(value) = value {
1264            if builder.validate_value(value).is_err() {
1265                return Err(SemanticAdError::ForeignValue { field, index });
1266            }
1267        }
1268    }
1269    Ok(())
1270}
1271
1272fn validate_len(
1273    field: &'static str,
1274    expected: usize,
1275    actual: usize,
1276) -> Result<(), SemanticAdError> {
1277    if expected != actual {
1278        return Err(SemanticAdError::Arity {
1279            field,
1280            expected,
1281            actual,
1282        });
1283    }
1284    Ok(())
1285}
1286
1287fn validate_insert<T>(
1288    map: &HashMap<&'static str, T>,
1289    family_id: &'static str,
1290    role: SemanticAdRuleRole,
1291) -> Result<(), SemanticExtensionRegistryError> {
1292    if !is_valid_family_id(family_id) {
1293        return Err(SemanticExtensionRegistryError::MalformedFamilyId { family_id });
1294    }
1295    if map.contains_key(family_id) {
1296        return Err(SemanticExtensionRegistryError::DuplicateRule { family_id, role });
1297    }
1298    Ok(())
1299}
1300
1301fn is_valid_family_id(family_id: &str) -> bool {
1302    let Some((prefix, version)) = family_id.rsplit_once('.') else {
1303        return false;
1304    };
1305    let Some(version) = version.strip_prefix('v') else {
1306        return false;
1307    };
1308    let Some((crate_name, op_name)) = prefix.split_once('.') else {
1309        return false;
1310    };
1311    !crate_name.is_empty()
1312        && !op_name.is_empty()
1313        && !version.is_empty()
1314        && version.bytes().all(|byte| byte.is_ascii_digit())
1315        && crate_name.is_ascii()
1316        && op_name.is_ascii()
1317        && !crate_name.chars().any(char::is_whitespace)
1318        && !op_name.chars().any(char::is_whitespace)
1319}
1320
1321#[cfg(test)]
1322mod tests;