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;