tenferro_linalg/ad/
support.rs

1#[cfg(test)]
2use crate::extension::LinalgOp;
3
4/// AD rule support status for a linalg operation or output.
5///
6/// # Examples
7///
8/// ```rust
9/// use tenferro_linalg::{linalg_ad_support, LinalgAdOpKind, LinalgAdRuleSupport};
10///
11/// let svd = linalg_ad_support(LinalgAdOpKind::Svd);
12/// assert_eq!(svd.linearize, LinalgAdRuleSupport::SupportedViaLinearize);
13/// ```
14#[derive(Clone, Copy, Debug, PartialEq, Eq)]
15pub enum LinalgAdRuleSupport {
16    Supported,
17    SupportedViaLinearize,
18    PartiallySupported,
19    NonDifferentiable,
20    Unsupported,
21    PendingOracle,
22}
23
24/// Implementation route used for a user-visible AD mode.
25///
26/// # Examples
27///
28/// ```rust
29/// use tenferro_linalg::{linalg_ad_support, LinalgAdOpKind, LinalgAdRoute};
30///
31/// let svd = linalg_ad_support(LinalgAdOpKind::Svd);
32/// assert_eq!(svd.vjp.route, LinalgAdRoute::LinearizeThenTranspose);
33/// ```
34#[derive(Clone, Copy, Debug, PartialEq, Eq)]
35pub enum LinalgAdRoute {
36    /// No supported route exists.
37    Unsupported,
38    /// The mode is emitted directly by the operation's linearize rule.
39    Linearize,
40    /// Reverse mode is supported by linearizing the primal graph and then
41    /// transposing that linear graph.
42    LinearizeThenTranspose,
43    /// Reverse mode is supported by linearizing the primal graph and using a
44    /// custom transposed-linear rule for the emitted linear operation.
45    LinearizeThenCustomLinearTranspose,
46    /// Reverse mode is emitted by a direct operation-specific primal VJP rule.
47    CustomVjp,
48    /// A custom VJP is preferred, with the canonical linearize-then-transpose
49    /// route retained as an intentional not-applicable fallback.
50    CustomPreferredWithLinearizeFallback,
51}
52
53/// User-visible AD support for one mode.
54///
55/// # Examples
56///
57/// ```rust
58/// use tenferro_linalg::{linalg_ad_support, LinalgAdOpKind, LinalgAdRoute};
59///
60/// let solve = linalg_ad_support(LinalgAdOpKind::TriangularSolve);
61/// assert_eq!(solve.vjp.route, LinalgAdRoute::LinearizeThenCustomLinearTranspose);
62/// ```
63#[derive(Clone, Copy, Debug, PartialEq, Eq)]
64pub struct LinalgAdModeSupport {
65    /// Whether the user-visible mode is supported.
66    pub status: LinalgAdRuleSupport,
67    /// How the mode is implemented.
68    pub route: LinalgAdRoute,
69}
70
71/// Operation keys covered by the linalg AD support manifest.
72///
73/// # Examples
74///
75/// ```rust
76/// use tenferro_linalg::LinalgAdOpKind;
77///
78/// assert!(LinalgAdOpKind::Svd.as_index() < LinalgAdOpKind::COUNT);
79/// ```
80#[derive(Clone, Copy, Debug, PartialEq, Eq)]
81pub enum LinalgAdOpKind {
82    Cholesky,
83    Lu,
84    LuFactor,
85    LuSolvePrepared,
86    FullPivLu,
87    FullPivLuSolve,
88    Svd,
89    SvdVals,
90    Qr,
91    Eigh,
92    EighVals,
93    Eig,
94    EigVals,
95    TriangularSolve,
96}
97
98impl LinalgAdOpKind {
99    pub const COUNT: usize = 14;
100
101    /// Return the manifest index for this operation kind.
102    ///
103    /// # Examples
104    ///
105    /// ```rust
106    /// use tenferro_linalg::LinalgAdOpKind;
107    ///
108    /// assert_eq!(LinalgAdOpKind::Cholesky.as_index(), 0);
109    /// ```
110    pub const fn as_index(self) -> usize {
111        match self {
112            Self::Cholesky => 0,
113            Self::Lu => 1,
114            Self::LuFactor => 2,
115            Self::LuSolvePrepared => 3,
116            Self::FullPivLu => 4,
117            Self::FullPivLuSolve => 5,
118            Self::Svd => 6,
119            Self::SvdVals => 7,
120            Self::Qr => 8,
121            Self::Eigh => 9,
122            Self::EighVals => 10,
123            Self::Eig => 11,
124            Self::EigVals => 12,
125            Self::TriangularSolve => 13,
126        }
127    }
128
129    #[cfg(test)]
130    pub(crate) const fn from_linalg_op(op: LinalgOp) -> Self {
131        match op {
132            LinalgOp::Cholesky => Self::Cholesky,
133            LinalgOp::Lu => Self::Lu,
134            LinalgOp::LuFactor => Self::LuFactor,
135            LinalgOp::LuSolvePrepared { .. } => Self::LuSolvePrepared,
136            LinalgOp::FullPivLu => Self::FullPivLu,
137            LinalgOp::FullPivLuSolve { .. } => Self::FullPivLuSolve,
138            LinalgOp::Svd { .. } => Self::Svd,
139            LinalgOp::SvdVals { .. } => Self::SvdVals,
140            LinalgOp::Qr { .. } => Self::Qr,
141            LinalgOp::Eigh { .. } => Self::Eigh,
142            LinalgOp::EighVals { .. } => Self::EighVals,
143            LinalgOp::Eig { .. } => Self::Eig,
144            LinalgOp::EigVals { .. } => Self::EigVals,
145            LinalgOp::TriangularSolve { .. } => Self::TriangularSolve,
146        }
147    }
148}
149
150/// AD support status for one output of a linalg operation.
151///
152/// # Examples
153///
154/// ```rust
155/// use tenferro_linalg::{linalg_ad_support, LinalgAdOpKind, LinalgAdRuleSupport};
156///
157/// let full_piv_lu = linalg_ad_support(LinalgAdOpKind::FullPivLu);
158/// let l_output = full_piv_lu.outputs.iter().find(|output| output.name == "l").unwrap();
159/// assert_eq!(l_output.status, LinalgAdRuleSupport::SupportedViaLinearize);
160/// ```
161#[derive(Clone, Copy, Debug, PartialEq, Eq)]
162pub struct LinalgAdOutputSupport {
163    /// Output position in the linalg operation result tuple.
164    pub index: usize,
165    /// Stable output name used by tests and support dashboards.
166    pub name: &'static str,
167    /// AD support status for this specific output.
168    pub status: LinalgAdRuleSupport,
169}
170
171/// AD support manifest entry for one linalg operation.
172///
173/// # Examples
174///
175/// ```rust
176/// use tenferro_linalg::{linalg_ad_support, LinalgAdOpKind, LinalgAdRuleSupport};
177///
178/// let solve = linalg_ad_support(LinalgAdOpKind::TriangularSolve);
179/// assert_eq!(solve.vjp.route, tenferro_linalg::LinalgAdRoute::LinearizeThenCustomLinearTranspose);
180/// ```
181#[derive(Clone, Copy, Debug, PartialEq, Eq)]
182pub struct LinalgAdSupport {
183    /// Operation kind described by this manifest entry.
184    pub kind: LinalgAdOpKind,
185    /// User-visible JVP support and route.
186    pub jvp: LinalgAdModeSupport,
187    /// User-visible VJP support and route.
188    pub vjp: LinalgAdModeSupport,
189    /// Definitional linearize rule implementation status.
190    pub linearize_rule: LinalgAdRuleSupport,
191    /// Direct primal VJP rule implementation status.
192    pub custom_vjp_rule: LinalgAdRuleSupport,
193    /// Custom transposed-linear rule implementation status.
194    pub custom_linear_transpose_rule: LinalgAdRuleSupport,
195    /// Forward-mode graph emission support.
196    pub linearize: LinalgAdRuleSupport,
197    /// Transposed-linear graph emission support.
198    pub transpose: LinalgAdRuleSupport,
199    /// Per-output support status for multi-output operations.
200    pub outputs: &'static [LinalgAdOutputSupport],
201    /// Numerical or semantic caveats for this operation family.
202    pub caveats: &'static [&'static str],
203}
204
205const fn mode(status: LinalgAdRuleSupport, route: LinalgAdRoute) -> LinalgAdModeSupport {
206    LinalgAdModeSupport { status, route }
207}
208
209const fn jvp_route(status: LinalgAdRuleSupport) -> LinalgAdRoute {
210    match status {
211        LinalgAdRuleSupport::Unsupported
212        | LinalgAdRuleSupport::NonDifferentiable
213        | LinalgAdRuleSupport::PendingOracle => LinalgAdRoute::Unsupported,
214        LinalgAdRuleSupport::Supported
215        | LinalgAdRuleSupport::SupportedViaLinearize
216        | LinalgAdRuleSupport::PartiallySupported => LinalgAdRoute::Linearize,
217    }
218}
219
220const fn support_entry(
221    kind: LinalgAdOpKind,
222    linearize: LinalgAdRuleSupport,
223    transpose: LinalgAdRuleSupport,
224    vjp: LinalgAdModeSupport,
225    custom_linear_transpose_rule: LinalgAdRuleSupport,
226    outputs: &'static [LinalgAdOutputSupport],
227    caveats: &'static [&'static str],
228) -> LinalgAdSupport {
229    LinalgAdSupport {
230        kind,
231        jvp: mode(linearize, jvp_route(linearize)),
232        vjp,
233        linearize_rule: linearize,
234        custom_vjp_rule: LinalgAdRuleSupport::Unsupported,
235        custom_linear_transpose_rule,
236        linearize,
237        transpose,
238        outputs,
239        caveats,
240    }
241}
242
243const fn output(
244    index: usize,
245    name: &'static str,
246    status: LinalgAdRuleSupport,
247) -> LinalgAdOutputSupport {
248    LinalgAdOutputSupport {
249        index,
250        name,
251        status,
252    }
253}
254
255static CHOLESKY_OUTPUTS: [LinalgAdOutputSupport; 1] = [output(
256    0,
257    "factor",
258    LinalgAdRuleSupport::SupportedViaLinearize,
259)];
260static LU_OUTPUTS: [LinalgAdOutputSupport; 4] = [
261    output(0, "p", LinalgAdRuleSupport::NonDifferentiable),
262    output(1, "l", LinalgAdRuleSupport::SupportedViaLinearize),
263    output(2, "u", LinalgAdRuleSupport::SupportedViaLinearize),
264    output(3, "parity", LinalgAdRuleSupport::NonDifferentiable),
265];
266static LU_FACTOR_OUTPUTS: [LinalgAdOutputSupport; 3] = [
267    output(0, "packed_lu", LinalgAdRuleSupport::Unsupported),
268    output(1, "pivots", LinalgAdRuleSupport::NonDifferentiable),
269    output(2, "parity", LinalgAdRuleSupport::NonDifferentiable),
270];
271static SOLUTION_OUTPUTS: [LinalgAdOutputSupport; 1] = [output(
272    0,
273    "solution",
274    LinalgAdRuleSupport::SupportedViaLinearize,
275)];
276static FULL_PIV_LU_OUTPUTS: [LinalgAdOutputSupport; 5] = [
277    output(0, "p", LinalgAdRuleSupport::NonDifferentiable),
278    output(1, "l", LinalgAdRuleSupport::SupportedViaLinearize),
279    output(2, "u", LinalgAdRuleSupport::SupportedViaLinearize),
280    output(3, "q", LinalgAdRuleSupport::NonDifferentiable),
281    output(4, "parity", LinalgAdRuleSupport::NonDifferentiable),
282];
283static FULL_PIV_LU_SOLVE_OUTPUTS: [LinalgAdOutputSupport; 1] = [output(
284    0,
285    "solution",
286    LinalgAdRuleSupport::SupportedViaLinearize,
287)];
288static SVD_OUTPUTS: [LinalgAdOutputSupport; 3] = [
289    output(0, "u", LinalgAdRuleSupport::SupportedViaLinearize),
290    output(
291        1,
292        "singular_values",
293        LinalgAdRuleSupport::SupportedViaLinearize,
294    ),
295    output(2, "vt", LinalgAdRuleSupport::SupportedViaLinearize),
296];
297static SVD_VALS_OUTPUTS: [LinalgAdOutputSupport; 1] = [output(
298    0,
299    "singular_values",
300    LinalgAdRuleSupport::SupportedViaLinearize,
301)];
302static QR_OUTPUTS: [LinalgAdOutputSupport; 2] = [
303    output(0, "q", LinalgAdRuleSupport::SupportedViaLinearize),
304    output(1, "r", LinalgAdRuleSupport::SupportedViaLinearize),
305];
306static EIGH_OUTPUTS: [LinalgAdOutputSupport; 2] = [
307    output(0, "eigenvalues", LinalgAdRuleSupport::SupportedViaLinearize),
308    output(
309        1,
310        "eigenvectors",
311        LinalgAdRuleSupport::SupportedViaLinearize,
312    ),
313];
314static EIGH_VALS_OUTPUTS: [LinalgAdOutputSupport; 1] = [output(
315    0,
316    "eigenvalues",
317    LinalgAdRuleSupport::SupportedViaLinearize,
318)];
319static EIG_OUTPUTS: [LinalgAdOutputSupport; 2] = [
320    output(0, "eigenvalues", LinalgAdRuleSupport::SupportedViaLinearize),
321    output(1, "eigenvectors", LinalgAdRuleSupport::Unsupported),
322];
323static EIG_VALS_OUTPUTS: [LinalgAdOutputSupport; 1] = [output(
324    0,
325    "eigenvalues",
326    LinalgAdRuleSupport::SupportedViaLinearize,
327)];
328
329static DECOMPOSITION_CAVEATS: [&str; 1] = [
330    "Derivative regularization handles near-degenerate spectra but does not make exact degeneracies smoothly differentiable.",
331];
332
333static LINALG_AD_SUPPORT: [LinalgAdSupport; LinalgAdOpKind::COUNT] = [
334    support_entry(
335        LinalgAdOpKind::Cholesky,
336        LinalgAdRuleSupport::SupportedViaLinearize,
337        LinalgAdRuleSupport::Unsupported,
338        mode(
339            LinalgAdRuleSupport::SupportedViaLinearize,
340            LinalgAdRoute::LinearizeThenTranspose,
341        ),
342        LinalgAdRuleSupport::Unsupported,
343        &CHOLESKY_OUTPUTS,
344        &DECOMPOSITION_CAVEATS,
345    ),
346    support_entry(
347        LinalgAdOpKind::Lu,
348        LinalgAdRuleSupport::PartiallySupported,
349        LinalgAdRuleSupport::Unsupported,
350        mode(
351            LinalgAdRuleSupport::PartiallySupported,
352            LinalgAdRoute::LinearizeThenTranspose,
353        ),
354        LinalgAdRuleSupport::Unsupported,
355        &LU_OUTPUTS,
356        &DECOMPOSITION_CAVEATS,
357    ),
358    support_entry(
359        LinalgAdOpKind::LuFactor,
360        LinalgAdRuleSupport::Unsupported,
361        LinalgAdRuleSupport::Unsupported,
362        mode(LinalgAdRuleSupport::Unsupported, LinalgAdRoute::Unsupported),
363        LinalgAdRuleSupport::Unsupported,
364        &LU_FACTOR_OUTPUTS,
365        &[],
366    ),
367    support_entry(
368        LinalgAdOpKind::LuSolvePrepared,
369        LinalgAdRuleSupport::SupportedViaLinearize,
370        LinalgAdRuleSupport::PartiallySupported,
371        mode(
372            LinalgAdRuleSupport::SupportedViaLinearize,
373            LinalgAdRoute::LinearizeThenCustomLinearTranspose,
374        ),
375        LinalgAdRuleSupport::PartiallySupported,
376        &SOLUTION_OUTPUTS,
377        &[],
378    ),
379    support_entry(
380        LinalgAdOpKind::FullPivLu,
381        LinalgAdRuleSupport::SupportedViaLinearize,
382        LinalgAdRuleSupport::Unsupported,
383        mode(
384            LinalgAdRuleSupport::SupportedViaLinearize,
385            LinalgAdRoute::LinearizeThenTranspose,
386        ),
387        LinalgAdRuleSupport::Unsupported,
388        &FULL_PIV_LU_OUTPUTS,
389        &DECOMPOSITION_CAVEATS,
390    ),
391    support_entry(
392        LinalgAdOpKind::FullPivLuSolve,
393        LinalgAdRuleSupport::SupportedViaLinearize,
394        LinalgAdRuleSupport::Supported,
395        mode(
396            LinalgAdRuleSupport::SupportedViaLinearize,
397            LinalgAdRoute::LinearizeThenCustomLinearTranspose,
398        ),
399        LinalgAdRuleSupport::Supported,
400        &FULL_PIV_LU_SOLVE_OUTPUTS,
401        &[],
402    ),
403    support_entry(
404        LinalgAdOpKind::Svd,
405        LinalgAdRuleSupport::SupportedViaLinearize,
406        LinalgAdRuleSupport::Unsupported,
407        mode(
408            LinalgAdRuleSupport::SupportedViaLinearize,
409            LinalgAdRoute::LinearizeThenTranspose,
410        ),
411        LinalgAdRuleSupport::Unsupported,
412        &SVD_OUTPUTS,
413        &DECOMPOSITION_CAVEATS,
414    ),
415    support_entry(
416        LinalgAdOpKind::SvdVals,
417        LinalgAdRuleSupport::SupportedViaLinearize,
418        LinalgAdRuleSupport::Unsupported,
419        mode(
420            LinalgAdRuleSupport::SupportedViaLinearize,
421            LinalgAdRoute::LinearizeThenTranspose,
422        ),
423        LinalgAdRuleSupport::Unsupported,
424        &SVD_VALS_OUTPUTS,
425        &DECOMPOSITION_CAVEATS,
426    ),
427    support_entry(
428        LinalgAdOpKind::Qr,
429        LinalgAdRuleSupport::SupportedViaLinearize,
430        LinalgAdRuleSupport::Unsupported,
431        mode(
432            LinalgAdRuleSupport::SupportedViaLinearize,
433            LinalgAdRoute::LinearizeThenTranspose,
434        ),
435        LinalgAdRuleSupport::Unsupported,
436        &QR_OUTPUTS,
437        &DECOMPOSITION_CAVEATS,
438    ),
439    support_entry(
440        LinalgAdOpKind::Eigh,
441        LinalgAdRuleSupport::SupportedViaLinearize,
442        LinalgAdRuleSupport::Unsupported,
443        mode(
444            LinalgAdRuleSupport::SupportedViaLinearize,
445            LinalgAdRoute::LinearizeThenTranspose,
446        ),
447        LinalgAdRuleSupport::Unsupported,
448        &EIGH_OUTPUTS,
449        &DECOMPOSITION_CAVEATS,
450    ),
451    support_entry(
452        LinalgAdOpKind::EighVals,
453        LinalgAdRuleSupport::SupportedViaLinearize,
454        LinalgAdRuleSupport::Unsupported,
455        mode(
456            LinalgAdRuleSupport::SupportedViaLinearize,
457            LinalgAdRoute::LinearizeThenTranspose,
458        ),
459        LinalgAdRuleSupport::Unsupported,
460        &EIGH_VALS_OUTPUTS,
461        &DECOMPOSITION_CAVEATS,
462    ),
463    support_entry(
464        LinalgAdOpKind::Eig,
465        LinalgAdRuleSupport::PartiallySupported,
466        LinalgAdRuleSupport::Unsupported,
467        mode(
468            LinalgAdRuleSupport::PartiallySupported,
469            LinalgAdRoute::LinearizeThenTranspose,
470        ),
471        LinalgAdRuleSupport::Unsupported,
472        &EIG_OUTPUTS,
473        &DECOMPOSITION_CAVEATS,
474    ),
475    support_entry(
476        LinalgAdOpKind::EigVals,
477        LinalgAdRuleSupport::SupportedViaLinearize,
478        LinalgAdRuleSupport::Unsupported,
479        mode(
480            LinalgAdRuleSupport::SupportedViaLinearize,
481            LinalgAdRoute::LinearizeThenTranspose,
482        ),
483        LinalgAdRuleSupport::Unsupported,
484        &EIG_VALS_OUTPUTS,
485        &DECOMPOSITION_CAVEATS,
486    ),
487    support_entry(
488        LinalgAdOpKind::TriangularSolve,
489        LinalgAdRuleSupport::SupportedViaLinearize,
490        LinalgAdRuleSupport::Supported,
491        mode(
492            LinalgAdRuleSupport::SupportedViaLinearize,
493            LinalgAdRoute::LinearizeThenCustomLinearTranspose,
494        ),
495        LinalgAdRuleSupport::Supported,
496        &SOLUTION_OUTPUTS,
497        &[],
498    ),
499];
500
501/// Return the complete linalg AD support manifest.
502///
503/// # Examples
504///
505/// ```rust
506/// let manifest = tenferro_linalg::all_linalg_ad_support();
507/// assert_eq!(manifest.len(), tenferro_linalg::LinalgAdOpKind::COUNT);
508/// ```
509pub fn all_linalg_ad_support() -> &'static [LinalgAdSupport; LinalgAdOpKind::COUNT] {
510    &LINALG_AD_SUPPORT
511}
512
513/// Return the support manifest entry for one linalg operation kind.
514///
515/// # Examples
516///
517/// ```rust
518/// use tenferro_linalg::{linalg_ad_support, LinalgAdOpKind};
519///
520/// let entry = linalg_ad_support(LinalgAdOpKind::Eigh);
521/// assert_eq!(entry.kind, LinalgAdOpKind::Eigh);
522/// ```
523pub fn linalg_ad_support(kind: LinalgAdOpKind) -> &'static LinalgAdSupport {
524    &LINALG_AD_SUPPORT[kind.as_index()]
525}
526
527#[cfg(test)]
528pub(crate) fn linalg_ad_support_for_op(op: LinalgOp) -> &'static LinalgAdSupport {
529    linalg_ad_support(LinalgAdOpKind::from_linalg_op(op))
530}
531
532#[cfg(test)]
533mod tests;
tenferro_linalg/ad/support.rs

tenferro_linalg/ad/
support.rs
