tenferro_linalg/
ad.rs

1//! Automatic differentiation support for `tenferro-linalg`.
2//!
3//! This module is enabled by the `autodiff` feature. It provides the linalg
4//! extension rule set used by explicit `tenferro_ad::AdContext` values.
5//!
6//! # Examples
7//!
8//! ```rust
9//! use tenferro_ad::AdContext;
10//! use tenferro_linalg::TracedTensorLinalgExt;
11//! use tenferro_runtime::TracedTensor;
12//!
13//! let ad = AdContext::builder()
14//!     .with_extension_rules(tenferro_linalg::ad_rules().unwrap())
15//!     .build()
16//!     .unwrap();
17//! let x = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0]).unwrap();
18//! let (_u, s, _vt) = x.svd().unwrap();
19//! let loss = s.reduce_sum(&[0]).unwrap();
20//! let grad = ad.grad(&loss, &x).unwrap();
21//! assert_eq!(grad.rank, 2);
22//! ```
23
24use std::sync::Arc;
25
26use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef};
27use tenferro_ad::extension::{
28    ExtensionLinearTransposeRule, ExtensionLinearizeRule, ExtensionOp, ExtensionRegistryError,
29    ExtensionRuleSet,
30};
31use tenferro_ops::ad::PrimitiveRuleBuilder;
32use tenferro_ops::std_tensor_op::StdTensorOp;
33use tenferro_ops::ShapeGuardContext;
34use tidu::{ADRuleError, ADRuleKind, ADRuleResult, PrimitiveTransposeInput};
35
36use crate::extension::{LinalgExtensionOp, LinalgOp};
37use crate::LINALG_EXTENSION_FAMILY_ID;
38
39mod rules;
40pub mod support;
41
42/// Return the explicit linalg extension AD rule set.
43///
44/// # Examples
45///
46/// ```rust
47/// let rules = tenferro_linalg::ad_rules().unwrap();
48/// assert!(rules.is_linearize_registered(tenferro_linalg::LINALG_EXTENSION_FAMILY_ID));
49/// assert!(rules.is_linear_transpose_registered(tenferro_linalg::LINALG_EXTENSION_FAMILY_ID));
50/// ```
51pub fn ad_rules() -> Result<ExtensionRuleSet, ExtensionRegistryError> {
52    ExtensionRuleSet::new()
53        .with_linearize(Arc::new(LinalgAdRule))?
54        .with_linear_transpose(Arc::new(LinalgAdRule))
55}
56
57#[derive(Debug)]
58struct LinalgAdRule;
59
60impl ExtensionLinearizeRule for LinalgAdRule {
61    fn family_id(&self) -> &'static str {
62        LINALG_EXTENSION_FAMILY_ID
63    }
64
65    fn linearize(
66        &self,
67        op: &dyn ExtensionOp,
68        builder: &mut dyn PrimitiveRuleBuilder,
69        primal_in: &[ValueKey<StdTensorOp>],
70        primal_out: &[ValueKey<StdTensorOp>],
71        tangent_in: &[Option<LocalValueId>],
72        ctx: &mut ShapeGuardContext,
73    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
74        let op = downcast_ad_op(op, ADRuleKind::Jvp)?;
75        match op.op() {
76            LinalgOp::Lu => rules::linearize_lu(builder, primal_in, primal_out, tangent_in, ctx),
77            LinalgOp::LuFactor => Ok(vec![None; op.output_count()]),
78            LinalgOp::LuSolvePrepared {
79                transpose_a,
80                conjugate_a,
81            } => rules::linearize_lu_solve_prepared(
82                builder,
83                primal_in,
84                primal_out,
85                tangent_in,
86                transpose_a,
87                conjugate_a,
88                ctx,
89            ),
90            LinalgOp::FullPivLu => {
91                rules::linearize_full_piv_lu(builder, primal_in, primal_out, tangent_in, ctx)
92            }
93            LinalgOp::FullPivLuSolve { transpose_a } => rules::linearize_full_piv_lu_solve(
94                builder,
95                primal_in,
96                primal_out,
97                tangent_in,
98                transpose_a,
99                ctx,
100            ),
101            LinalgOp::TriangularSolve {
102                left_side,
103                lower,
104                transpose_a,
105                unit_diagonal,
106            } => rules::linearize_triangular_solve(
107                builder,
108                primal_in,
109                primal_out,
110                tangent_in,
111                rules::TriangularSolveFlags::new(left_side, lower, transpose_a, unit_diagonal),
112                ctx,
113            ),
114            LinalgOp::Cholesky => {
115                rules::linearize_cholesky(builder, primal_in, primal_out, tangent_in, ctx)
116            }
117            LinalgOp::Svd { derivative_eps, .. } => rules::linearize_svd(
118                builder,
119                primal_in,
120                primal_out,
121                tangent_in,
122                derivative_eps,
123                ctx,
124            ),
125            LinalgOp::SvdVals { derivative_eps } => {
126                rules::linearize_svd_values(builder, primal_in, tangent_in, derivative_eps, ctx)
127            }
128            LinalgOp::Qr { .. } => {
129                rules::linearize_qr(builder, primal_in, primal_out, tangent_in, ctx)
130            }
131            LinalgOp::Eigh { derivative_eps, .. } => rules::linearize_eigh(
132                builder,
133                primal_in,
134                primal_out,
135                tangent_in,
136                derivative_eps,
137                ctx,
138            ),
139            LinalgOp::EighVals { derivative_eps } => {
140                rules::linearize_eigh_values(builder, primal_in, tangent_in, derivative_eps, ctx)
141            }
142            LinalgOp::Eig { input_dtype } => {
143                rules::linearize_eig(builder, primal_in, primal_out, tangent_in, input_dtype, ctx)
144            }
145            LinalgOp::EigVals { input_dtype } => {
146                rules::linearize_eig_values(builder, primal_in, tangent_in, input_dtype, ctx)
147            }
148        }
149    }
150}
151
152impl ExtensionLinearTransposeRule for LinalgAdRule {
153    fn family_id(&self) -> &'static str {
154        LINALG_EXTENSION_FAMILY_ID
155    }
156
157    fn linear_transpose(
158        &self,
159        op: &dyn ExtensionOp,
160        builder: &mut dyn PrimitiveRuleBuilder,
161        cotangent_out: &[Option<LocalValueId>],
162        inputs: &[PrimitiveTransposeInput<StdTensorOp>],
163        active_mask: &[bool],
164        ctx: &mut ShapeGuardContext,
165    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
166        let op = downcast_ad_op(op, ADRuleKind::Transpose)?;
167        let mut builder = DynBuilder(builder);
168        let mode = OperationRole::Linearized {
169            active_mask: active_mask.to_vec(),
170        };
171        match op.op() {
172            LinalgOp::TriangularSolve {
173                left_side,
174                lower,
175                transpose_a,
176                unit_diagonal,
177            } => {
178                let value_inputs =
179                    linear_solve_transpose_inputs("triangular_solve", inputs, active_mask)?;
180                rules::transpose_triangular_solve(
181                    &mut builder,
182                    cotangent_out,
183                    &value_inputs,
184                    &mode,
185                    rules::TriangularSolveFlags::new(left_side, lower, transpose_a, unit_diagonal),
186                    ctx,
187                )
188            }
189            LinalgOp::LuSolvePrepared {
190                transpose_a,
191                conjugate_a,
192            } => {
193                let value_inputs = lu_solve_prepared_transpose_inputs(inputs, active_mask)?;
194                rules::transpose_lu_solve_prepared(
195                    &mut builder,
196                    cotangent_out,
197                    &value_inputs,
198                    &mode,
199                    transpose_a,
200                    conjugate_a,
201                    ctx,
202                )
203            }
204            LinalgOp::FullPivLuSolve { transpose_a } => {
205                let value_inputs =
206                    linear_solve_transpose_inputs("full_piv_lu_solve", inputs, active_mask)?;
207                rules::transpose_full_piv_lu_solve(
208                    &mut builder,
209                    cotangent_out,
210                    &value_inputs,
211                    &mode,
212                    transpose_a,
213                    ctx,
214                )
215            }
216            LinalgOp::Cholesky
217            | LinalgOp::Lu
218            | LinalgOp::LuFactor
219            | LinalgOp::FullPivLu
220            | LinalgOp::Svd { .. }
221            | LinalgOp::SvdVals { .. }
222            | LinalgOp::Qr { .. }
223            | LinalgOp::Eigh { .. }
224            | LinalgOp::EighVals { .. }
225            | LinalgOp::Eig { .. }
226            | LinalgOp::EigVals { .. } => Ok(vec![None; op.input_count()]),
227        }
228    }
229}
230
231struct DynBuilder<'a>(&'a mut dyn PrimitiveRuleBuilder);
232
233impl PrimitiveRuleBuilder for DynBuilder<'_> {
234    fn add_operation(
235        &mut self,
236        op: StdTensorOp,
237        inputs: Vec<ValueRef<StdTensorOp>>,
238        mode: OperationRole,
239    ) -> Vec<LocalValueId> {
240        self.0.add_operation(op, inputs, mode)
241    }
242}
243
244fn downcast_ad_op(op: &dyn ExtensionOp, kind: ADRuleKind) -> ADRuleResult<&LinalgExtensionOp> {
245    op.as_any()
246        .downcast_ref::<LinalgExtensionOp>()
247        .ok_or_else(|| {
248            ADRuleError::invalid_input("tenferro-linalg.linalg.v1", kind, "payload type mismatch")
249        })
250}
251
252fn linear_solve_transpose_inputs(
253    op: &str,
254    inputs: &[PrimitiveTransposeInput<StdTensorOp>],
255    active_mask: &[bool],
256) -> ADRuleResult<Vec<ValueRef<StdTensorOp>>> {
257    let matrix_active = active_mask.first().copied().unwrap_or(false);
258    inputs
259        .iter()
260        .enumerate()
261        .map(|(index, input)| {
262            if index == 0 || matrix_active {
263                fixed_transpose_value(op, index, input)
264            } else {
265                Ok(metadata_transpose_value(input))
266            }
267        })
268        .collect()
269}
270
271fn lu_solve_prepared_transpose_inputs(
272    inputs: &[PrimitiveTransposeInput<StdTensorOp>],
273    active_mask: &[bool],
274) -> ADRuleResult<Vec<ValueRef<StdTensorOp>>> {
275    let matrix_active = active_mask.first().copied().unwrap_or(false);
276    inputs
277        .iter()
278        .enumerate()
279        .map(|(index, input)| {
280            if index <= 2 || matrix_active {
281                fixed_transpose_value("lu_solve_prepared", index, input)
282            } else {
283                Ok(metadata_transpose_value(input))
284            }
285        })
286        .collect()
287}
288
289fn metadata_transpose_value(input: &PrimitiveTransposeInput<StdTensorOp>) -> ValueRef<StdTensorOp> {
290    ValueRef::External(input.key().clone())
291}
292
293fn fixed_transpose_value(
294    op: &str,
295    index: usize,
296    input: &PrimitiveTransposeInput<StdTensorOp>,
297) -> ADRuleResult<ValueRef<StdTensorOp>> {
298    match input {
299        PrimitiveTransposeInput::Residual(key) => Ok(ValueRef::External(key.clone())),
300        PrimitiveTransposeInput::Linear {
301            primal: Some(primal),
302            ..
303        } => Ok(ValueRef::External(primal.clone())),
304        PrimitiveTransposeInput::Linear { key, primal: None } => {
305            Err(ADRuleError::invalid_input(
306                op,
307                ADRuleKind::Transpose,
308                format!(
309                    "transpose input {index} is linear-only and cannot be retained as a tensor operand: {key:?}"
310                ),
311            ))
312        }
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319    use crate::extension::{EighGauge, QrGauge, SvdGauge, DEFAULT_DECOMPOSITION_DERIVATIVE_EPS};
320    use computegraph::graph::GraphBuilder;
321    use std::collections::HashSet;
322    use tenferro_ops::input_key::TensorInputKey;
323    use tenferro_ops::{ShapeExtent, SymDim, TensorMeta};
324    use tenferro_tensor::DType;
325
326    fn input_key(id: u64) -> ValueKey<StdTensorOp> {
327        ValueKey::Input(TensorInputKey::User { id })
328    }
329
330    fn insert_meta(ctx: &mut ShapeGuardContext, key: ValueKey<StdTensorOp>, shape: &[usize]) {
331        ctx.insert_metadata(
332            key,
333            TensorMeta::exact(
334                DType::F64,
335                shape.iter().copied().map(SymDim::from).collect(),
336            ),
337        );
338    }
339
340    fn insert_typed_meta(
341        ctx: &mut ShapeGuardContext,
342        key: ValueKey<StdTensorOp>,
343        dtype: DType,
344        shape: &[usize],
345    ) {
346        ctx.insert_metadata(
347            key,
348            TensorMeta::exact(dtype, shape.iter().copied().map(SymDim::from).collect()),
349        );
350    }
351
352    fn eigh_context() -> (
353        ShapeGuardContext,
354        ValueKey<StdTensorOp>,
355        Vec<ValueKey<StdTensorOp>>,
356    ) {
357        let mut ctx = ShapeGuardContext::default();
358        let a = input_key(1);
359        let w = input_key(2);
360        let v = input_key(3);
361        insert_typed_meta(&mut ctx, a.clone(), DType::F64, &[2, 2]);
362        insert_typed_meta(&mut ctx, w.clone(), DType::F64, &[2]);
363        insert_typed_meta(&mut ctx, v.clone(), DType::F64, &[2, 2]);
364        (ctx, a, vec![w, v])
365    }
366
367    fn eig_context() -> (
368        ShapeGuardContext,
369        ValueKey<StdTensorOp>,
370        Vec<ValueKey<StdTensorOp>>,
371    ) {
372        let mut ctx = ShapeGuardContext::default();
373        let a = input_key(114);
374        let w = input_key(115);
375        let v = input_key(116);
376        insert_typed_meta(&mut ctx, a.clone(), DType::F64, &[2, 2]);
377        insert_typed_meta(&mut ctx, w.clone(), DType::C64, &[2]);
378        insert_typed_meta(&mut ctx, v.clone(), DType::C64, &[2, 2]);
379        (ctx, a, vec![w, v])
380    }
381
382    fn lu_context(
383        shape: &[usize],
384    ) -> (
385        ShapeGuardContext,
386        ValueKey<StdTensorOp>,
387        Vec<ValueKey<StdTensorOp>>,
388    ) {
389        let mut ctx = ShapeGuardContext::default();
390        let a = input_key(4);
391        let p = input_key(5);
392        let l = input_key(6);
393        let u = input_key(7);
394        let parity = input_key(8);
395        let k = shape[0].min(shape[1]);
396        insert_typed_meta(&mut ctx, a.clone(), DType::F64, shape);
397        insert_typed_meta(&mut ctx, p.clone(), DType::F64, &[shape[0], shape[0]]);
398        insert_typed_meta(&mut ctx, l.clone(), DType::F64, &[shape[0], k]);
399        insert_typed_meta(&mut ctx, u.clone(), DType::F64, &[k, shape[1]]);
400        insert_typed_meta(&mut ctx, parity.clone(), DType::F64, &[]);
401        (ctx, a, vec![p, l, u, parity])
402    }
403
404    fn svd_context(
405        shape: &[usize],
406    ) -> (
407        ShapeGuardContext,
408        ValueKey<StdTensorOp>,
409        Vec<ValueKey<StdTensorOp>>,
410    ) {
411        let mut ctx = ShapeGuardContext::default();
412        let a = input_key(120);
413        let u = input_key(121);
414        let s = input_key(122);
415        let vt = input_key(123);
416        let k = shape[0].min(shape[1]);
417        insert_typed_meta(&mut ctx, a.clone(), DType::F64, shape);
418        insert_typed_meta(&mut ctx, u.clone(), DType::F64, &[shape[0], k]);
419        insert_typed_meta(&mut ctx, s.clone(), DType::F64, &[k]);
420        insert_typed_meta(&mut ctx, vt.clone(), DType::F64, &[k, shape[1]]);
421        (ctx, a, vec![u, s, vt])
422    }
423
424    fn qr_context(
425        shape: &[usize],
426    ) -> (
427        ShapeGuardContext,
428        ValueKey<StdTensorOp>,
429        Vec<ValueKey<StdTensorOp>>,
430    ) {
431        let mut ctx = ShapeGuardContext::default();
432        let a = input_key(9);
433        let q = input_key(10);
434        let r = input_key(11);
435        let k = shape[0].min(shape[1]);
436        insert_typed_meta(&mut ctx, a.clone(), DType::F64, shape);
437        insert_typed_meta(&mut ctx, q.clone(), DType::F64, &[shape[0], k]);
438        insert_typed_meta(&mut ctx, r.clone(), DType::F64, &[k, shape[1]]);
439        (ctx, a, vec![q, r])
440    }
441
442    fn with_active_values(
443        ctx: ShapeGuardContext,
444        values: impl IntoIterator<Item = ValueKey<StdTensorOp>>,
445    ) -> ShapeGuardContext {
446        ctx.with_linearize_active_values(Arc::new(values.into_iter().collect::<HashSet<_>>()))
447    }
448
449    #[test]
450    fn full_piv_lu_jvp_returns_inactive_outputs_for_non_square_input() {
451        let mut builder = GraphBuilder::<StdTensorOp>::new();
452        let mut ctx = ShapeGuardContext::default();
453        let primal = input_key(1);
454        insert_meta(&mut ctx, primal.clone(), &[2, 3]);
455        let tangent = builder.add_input(TensorInputKey::User { id: 2 });
456        let outputs = [
457            input_key(10),
458            input_key(11),
459            input_key(12),
460            input_key(13),
461            input_key(14),
462        ];
463        let op = LinalgExtensionOp::new(LinalgOp::FullPivLu);
464
465        let result = LinalgAdRule
466            .linearize(
467                &op,
468                &mut builder,
469                &[primal],
470                &outputs,
471                &[Some(tangent)],
472                &mut ctx,
473            )
474            .unwrap();
475
476        assert_eq!(result, vec![None, None, None, None, None]);
477        assert!(builder.build().operations().is_empty());
478    }
479
480    #[test]
481    fn lu_linearize_prunes_inactive_factor_outputs() {
482        let op = LinalgExtensionOp::new(LinalgOp::Lu);
483        for (case, active_slot, expected_active) in [
484            ("l only", 1_usize, vec![false, true, false, false]),
485            ("u only", 2_usize, vec![false, false, true, false]),
486        ] {
487            let (ctx, a, outputs) = lu_context(&[2, 2]);
488            let mut ctx = with_active_values(ctx, [outputs[active_slot].clone()]);
489            let mut builder = GraphBuilder::<StdTensorOp>::new();
490            let tangent = builder.add_input(TensorInputKey::User { id: 130 });
491
492            let result = LinalgAdRule
493                .linearize(
494                    &op,
495                    &mut builder,
496                    &[a],
497                    &outputs,
498                    &[Some(tangent)],
499                    &mut ctx,
500                )
501                .unwrap();
502
503            assert_eq!(
504                result.iter().map(Option::is_some).collect::<Vec<_>>(),
505                expected_active,
506                "{case}"
507            );
508            let pruned_count = builder.build().operations().len();
509
510            let (full_ctx, full_a, full_outputs) = lu_context(&[2, 2]);
511            let mut full_ctx =
512                with_active_values(full_ctx, [full_outputs[1].clone(), full_outputs[2].clone()]);
513            let mut full_builder = GraphBuilder::<StdTensorOp>::new();
514            let full_tangent = full_builder.add_input(TensorInputKey::User { id: 131 });
515            let full_result = LinalgAdRule
516                .linearize(
517                    &op,
518                    &mut full_builder,
519                    &[full_a],
520                    &full_outputs,
521                    &[Some(full_tangent)],
522                    &mut full_ctx,
523                )
524                .unwrap();
525
526            assert_eq!(
527                full_result.iter().map(Option::is_some).collect::<Vec<_>>(),
528                vec![false, true, true, false],
529                "{case}"
530            );
531            let full_count = full_builder.build().operations().len();
532            assert!(
533                pruned_count < full_count,
534                "{case} should not emit both LU factor tangent branches: {pruned_count} >= {full_count}"
535            );
536        }
537    }
538
539    #[test]
540    fn one_input_linalg_jvps_prune_when_all_outputs_are_inactive() {
541        let cases = [
542            (
543                LinalgOp::Lu,
544                lu_context(&[2, 2]),
545                vec![None, None, None, None],
546            ),
547            (
548                LinalgOp::Svd {
549                    derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
550                    gauge: SvdGauge::Raw,
551                },
552                svd_context(&[2, 2]),
553                vec![None, None, None],
554            ),
555            (
556                LinalgOp::Eigh {
557                    derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
558                    gauge: EighGauge::Raw,
559                },
560                eigh_context(),
561                vec![None, None],
562            ),
563            (
564                LinalgOp::Eig {
565                    input_dtype: DType::F64,
566                },
567                eig_context(),
568                vec![None, None],
569            ),
570            (
571                LinalgOp::Qr {
572                    gauge: QrGauge::Raw,
573                },
574                qr_context(&[3, 2]),
575                vec![None, None],
576            ),
577        ];
578
579        for (kind, (ctx, a, outputs), expected) in cases {
580            let mut ctx = with_active_values(ctx, []);
581            let mut builder = GraphBuilder::<StdTensorOp>::new();
582            let tangent = builder.add_input(TensorInputKey::User { id: 132 });
583            let op = LinalgExtensionOp::new(kind);
584
585            let result = LinalgAdRule
586                .linearize(
587                    &op,
588                    &mut builder,
589                    &[a],
590                    &outputs,
591                    &[Some(tangent)],
592                    &mut ctx,
593                )
594                .unwrap();
595
596            assert_eq!(result, expected, "{kind:?}");
597            assert!(
598                builder.build().operations().is_empty(),
599                "{kind:?} should not emit AD graph operations for inactive outputs"
600            );
601        }
602    }
603
604    #[test]
605    fn svd_linearize_prunes_inactive_vector_outputs() {
606        let (ctx, a, outputs) = svd_context(&[2, 2]);
607        let mut ctx = with_active_values(ctx, [outputs[1].clone()]);
608        let mut builder = GraphBuilder::<StdTensorOp>::new();
609        let tangent = builder.add_input(TensorInputKey::User { id: 133 });
610        let op = LinalgExtensionOp::new(LinalgOp::Svd {
611            derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
612            gauge: SvdGauge::Raw,
613        });
614
615        let result = LinalgAdRule
616            .linearize(
617                &op,
618                &mut builder,
619                &[a],
620                &outputs,
621                &[Some(tangent)],
622                &mut ctx,
623            )
624            .unwrap();
625
626        assert_eq!(
627            result.iter().map(Option::is_some).collect::<Vec<_>>(),
628            vec![false, true, false]
629        );
630        assert!(
631            builder.build().operations().len() <= 5,
632            "singular-value-only SVD JVP should not emit the vector F-matrix chain"
633        );
634    }
635
636    #[test]
637    fn eigh_linearize_prunes_inactive_eigenvalue_output() {
638        let (ctx, a, outputs) = eigh_context();
639        let mut ctx = with_active_values(ctx, [outputs[1].clone()]);
640        let mut builder = GraphBuilder::<StdTensorOp>::new();
641        let tangent = builder.add_input(TensorInputKey::User { id: 134 });
642        let op = LinalgExtensionOp::new(LinalgOp::Eigh {
643            derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
644            gauge: EighGauge::Raw,
645        });
646
647        let result = LinalgAdRule
648            .linearize(
649                &op,
650                &mut builder,
651                &[a],
652                &outputs,
653                &[Some(tangent)],
654                &mut ctx,
655            )
656            .unwrap();
657
658        assert_eq!(
659            result.iter().map(Option::is_some).collect::<Vec<_>>(),
660            vec![false, true]
661        );
662    }
663
664    #[test]
665    fn eig_linearize_prunes_unsupported_inactive_eigenvalue_output() {
666        let (ctx, a, outputs) = eig_context();
667        let mut ctx = with_active_values(ctx, [outputs[1].clone()]);
668        let mut builder = GraphBuilder::<StdTensorOp>::new();
669        let tangent = builder.add_input(TensorInputKey::User { id: 137 });
670        let op = LinalgExtensionOp::new(LinalgOp::Eig {
671            input_dtype: DType::F64,
672        });
673
674        let result = LinalgAdRule
675            .linearize(
676                &op,
677                &mut builder,
678                &[a],
679                &outputs,
680                &[Some(tangent)],
681                &mut ctx,
682            )
683            .unwrap();
684
685        assert_eq!(result, vec![None, None]);
686        assert!(
687            builder.build().operations().is_empty(),
688            "eigenvectors-only Eig JVP is unsupported and should not emit eigenvalue tangent work"
689        );
690    }
691
692    #[test]
693    fn qr_linearize_prunes_inactive_factor_outputs() {
694        let op = LinalgExtensionOp::new(LinalgOp::Qr {
695            gauge: QrGauge::Raw,
696        });
697        for (case, active_slot, expected_active) in [
698            ("q only", 0_usize, vec![true, false]),
699            ("r only", 1_usize, vec![false, true]),
700        ] {
701            let (ctx, a, outputs) = qr_context(&[3, 2]);
702            let mut ctx = with_active_values(ctx, [outputs[active_slot].clone()]);
703            let mut builder = GraphBuilder::<StdTensorOp>::new();
704            let tangent = builder.add_input(TensorInputKey::User { id: 135 });
705
706            let result = LinalgAdRule
707                .linearize(
708                    &op,
709                    &mut builder,
710                    &[a],
711                    &outputs,
712                    &[Some(tangent)],
713                    &mut ctx,
714                )
715                .unwrap();
716
717            assert_eq!(
718                result.iter().map(Option::is_some).collect::<Vec<_>>(),
719                expected_active,
720                "{case}"
721            );
722            let pruned_count = builder.build().operations().len();
723
724            let (full_ctx, full_a, full_outputs) = qr_context(&[3, 2]);
725            let mut full_ctx =
726                with_active_values(full_ctx, [full_outputs[0].clone(), full_outputs[1].clone()]);
727            let mut full_builder = GraphBuilder::<StdTensorOp>::new();
728            let full_tangent = full_builder.add_input(TensorInputKey::User { id: 136 });
729            let full_result = LinalgAdRule
730                .linearize(
731                    &op,
732                    &mut full_builder,
733                    &[full_a],
734                    &full_outputs,
735                    &[Some(full_tangent)],
736                    &mut full_ctx,
737                )
738                .unwrap();
739
740            assert_eq!(
741                full_result.iter().map(Option::is_some).collect::<Vec<_>>(),
742                vec![true, true],
743                "{case}"
744            );
745            let full_count = full_builder.build().operations().len();
746            assert!(
747                pruned_count < full_count,
748                "{case} should not emit both QR factor tangent branches: {pruned_count} >= {full_count}"
749            );
750        }
751    }
752
753    #[test]
754    fn triangular_solve_jvp_rejects_non_matrix_operands() {
755        let mut builder = GraphBuilder::<StdTensorOp>::new();
756        let mut ctx = ShapeGuardContext::default();
757        let lhs = input_key(20);
758        let rhs = input_key(21);
759        insert_meta(&mut ctx, lhs.clone(), &[2, 2]);
760        insert_meta(&mut ctx, rhs.clone(), &[2]);
761        let rhs_tangent = builder.add_input(TensorInputKey::User { id: 22 });
762        let op = LinalgExtensionOp::new(LinalgOp::TriangularSolve {
763            left_side: true,
764            lower: true,
765            transpose_a: false,
766            unit_diagonal: false,
767        });
768
769        let err = LinalgAdRule
770            .linearize(
771                &op,
772                &mut builder,
773                &[lhs, rhs],
774                &[input_key(23)],
775                &[None, Some(rhs_tangent)],
776                &mut ctx,
777            )
778            .unwrap_err();
779
780        assert_eq!(err.rule(), ADRuleKind::Jvp);
781        assert!(err
782            .to_string()
783            .contains("expected matrix operands with rank >= 2"));
784        assert!(builder.build().operations().is_empty());
785    }
786
787    #[test]
788    fn triangular_solve_jvp_accepts_upper_bound_matrix_metadata() {
789        let mut builder = GraphBuilder::<StdTensorOp>::new();
790        let mut ctx = ShapeGuardContext::default();
791        let lhs = input_key(30);
792        let rhs = input_key(31);
793        ctx.insert_metadata(
794            lhs.clone(),
795            TensorMeta::with_extents(
796                DType::F64,
797                vec![
798                    ShapeExtent::upper_bound(SymDim::from(4usize)),
799                    ShapeExtent::upper_bound(SymDim::from(4usize)),
800                ],
801            ),
802        );
803        ctx.insert_metadata(
804            rhs.clone(),
805            TensorMeta::with_extents(
806                DType::F64,
807                vec![
808                    ShapeExtent::upper_bound(SymDim::from(4usize)),
809                    ShapeExtent::upper_bound(SymDim::from(2usize)),
810                ],
811            ),
812        );
813        let rhs_tangent = builder.add_input(TensorInputKey::User { id: 32 });
814        let op = LinalgExtensionOp::new(LinalgOp::TriangularSolve {
815            left_side: true,
816            lower: true,
817            transpose_a: false,
818            unit_diagonal: false,
819        });
820
821        let result = LinalgAdRule
822            .linearize(
823                &op,
824                &mut builder,
825                &[lhs.clone(), rhs],
826                &[input_key(33)],
827                &[None, Some(rhs_tangent)],
828                &mut ctx,
829            )
830            .unwrap();
831
832        assert!(result[0].is_some());
833        let graph = builder.build();
834        assert_eq!(graph.operations().len(), 1);
835        let solve = &graph.operations()[0];
836        assert_eq!(solve.inputs[0], ValueRef::External(lhs));
837        assert_eq!(solve.inputs[1], ValueRef::Local(rhs_tangent));
838    }
839
840    #[test]
841    fn triangular_solve_transpose_accepts_upper_bound_matrix_metadata() {
842        let mut builder = GraphBuilder::<StdTensorOp>::new();
843        let mut ctx = ShapeGuardContext::default();
844        let lhs = input_key(40);
845        let rhs = input_key(41);
846        ctx.insert_metadata(
847            lhs.clone(),
848            TensorMeta::with_extents(
849                DType::F64,
850                vec![
851                    ShapeExtent::upper_bound(SymDim::from(4usize)),
852                    ShapeExtent::upper_bound(SymDim::from(4usize)),
853                ],
854            ),
855        );
856        ctx.insert_metadata(
857            rhs.clone(),
858            TensorMeta::with_extents(
859                DType::F64,
860                vec![
861                    ShapeExtent::upper_bound(SymDim::from(4usize)),
862                    ShapeExtent::upper_bound(SymDim::from(2usize)),
863                ],
864            ),
865        );
866        let cotangent = builder.add_input(TensorInputKey::User { id: 42 });
867        let op = LinalgExtensionOp::new(LinalgOp::TriangularSolve {
868            left_side: true,
869            lower: true,
870            transpose_a: false,
871            unit_diagonal: false,
872        });
873
874        let result = LinalgAdRule
875            .linear_transpose(
876                &op,
877                &mut builder,
878                &[Some(cotangent)],
879                &[
880                    PrimitiveTransposeInput::Residual(lhs.clone()),
881                    PrimitiveTransposeInput::Residual(rhs),
882                ],
883                &[false, true],
884                &mut ctx,
885            )
886            .unwrap();
887
888        assert_eq!(result[0], None);
889        assert!(result[1].is_some());
890        let graph = builder.build();
891        assert_eq!(graph.operations().len(), 1);
892        assert_eq!(graph.operations()[0].inputs[0], ValueRef::External(lhs));
893        assert_eq!(graph.operations()[0].inputs[1], ValueRef::Local(cotangent));
894    }
895
896    #[test]
897    fn eigh_values_has_no_handwritten_direct_transpose() {
898        let (mut ctx, a, _primal_outputs) = eigh_context();
899        let mut builder = GraphBuilder::<StdTensorOp>::new();
900        let cotangent = builder.add_input(TensorInputKey::User { id: 85 });
901        let op = LinalgExtensionOp::new(LinalgOp::EighVals {
902            derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
903        });
904
905        let result = LinalgAdRule
906            .linear_transpose(
907                &op,
908                &mut builder,
909                &[Some(cotangent)],
910                &[PrimitiveTransposeInput::Residual(a)],
911                &[true],
912                &mut ctx,
913            )
914            .unwrap();
915
916        assert_eq!(result, vec![None]);
917        assert!(
918            builder.build().operations().is_empty(),
919            "EighVals reverse support should come from linearize + generic transpose"
920        );
921    }
922
923    #[test]
924    fn full_eigh_has_no_handwritten_direct_transpose() {
925        let (mut ctx, a, _primal_outputs) = eigh_context();
926        let mut builder = GraphBuilder::<StdTensorOp>::new();
927        let g_w = builder.add_input(TensorInputKey::User { id: 86 });
928        let g_v = builder.add_input(TensorInputKey::User { id: 87 });
929        let op = LinalgExtensionOp::new(LinalgOp::Eigh {
930            derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
931            gauge: EighGauge::Raw,
932        });
933
934        let result = LinalgAdRule
935            .linear_transpose(
936                &op,
937                &mut builder,
938                &[Some(g_w), Some(g_v)],
939                &[PrimitiveTransposeInput::Residual(a)],
940                &[true],
941                &mut ctx,
942            )
943            .unwrap();
944
945        assert_eq!(result, vec![None]);
946        assert!(
947            builder.build().operations().is_empty(),
948            "Eigh reverse support should come from linearize + generic transpose"
949        );
950    }
951
952    #[test]
953    fn qr_has_no_handwritten_direct_transpose() {
954        let (mut ctx, a, _primal_outputs) = qr_context(&[3, 2]);
955        let mut builder = GraphBuilder::<StdTensorOp>::new();
956        let g_q = builder.add_input(TensorInputKey::User { id: 88 });
957        let g_r = builder.add_input(TensorInputKey::User { id: 89 });
958
959        let result = LinalgAdRule
960            .linear_transpose(
961                &LinalgExtensionOp::new(LinalgOp::Qr {
962                    gauge: QrGauge::Raw,
963                }),
964                &mut builder,
965                &[Some(g_q), Some(g_r)],
966                &[PrimitiveTransposeInput::Residual(a)],
967                &[true],
968                &mut ctx,
969            )
970            .unwrap();
971
972        assert_eq!(result, vec![None]);
973        assert!(
974            builder.build().operations().is_empty(),
975            "QR reverse support should come from linearize + generic transpose"
976        );
977    }
978
979    #[test]
980    fn cholesky_jvp_uses_rank_when_input_metadata_is_upper_bound() {
981        let mut builder = GraphBuilder::<StdTensorOp>::new();
982        let mut ctx = ShapeGuardContext::default();
983        let primal = input_key(50);
984        ctx.insert_metadata(
985            primal.clone(),
986            TensorMeta::with_extents(
987                DType::F64,
988                vec![
989                    ShapeExtent::upper_bound(SymDim::from(4usize)),
990                    ShapeExtent::upper_bound(SymDim::from(4usize)),
991                ],
992            ),
993        );
994        let tangent = builder.add_input(TensorInputKey::User { id: 51 });
995        let op = LinalgExtensionOp::new(LinalgOp::Cholesky);
996
997        let result = LinalgAdRule
998            .linearize(
999                &op,
1000                &mut builder,
1001                &[primal],
1002                &[input_key(52)],
1003                &[Some(tangent)],
1004                &mut ctx,
1005            )
1006            .unwrap();
1007
1008        assert!(result[0].is_some());
1009        assert!(!builder.build().operations().is_empty());
1010    }
1011
1012    #[test]
1013    fn cholesky_jvp_propagates_missing_input_metadata() {
1014        let mut builder = GraphBuilder::<StdTensorOp>::new();
1015        let mut ctx = ShapeGuardContext::default();
1016        let primal = input_key(55);
1017        let tangent = builder.add_input(TensorInputKey::User { id: 56 });
1018        let op = LinalgExtensionOp::new(LinalgOp::Cholesky);
1019
1020        let err = LinalgAdRule
1021            .linearize(
1022                &op,
1023                &mut builder,
1024                &[primal],
1025                &[input_key(57)],
1026                &[Some(tangent)],
1027                &mut ctx,
1028            )
1029            .unwrap_err();
1030
1031        assert_eq!(err.rule(), ADRuleKind::Jvp);
1032        assert!(err.to_string().contains("missing TensorMeta"));
1033        assert!(builder.build().operations().is_empty());
1034    }
1035
1036    #[test]
1037    fn one_input_linalg_jvps_return_inactive_for_non_matrix_input() {
1038        let cases = [
1039            LinalgOp::Cholesky,
1040            LinalgOp::Lu,
1041            LinalgOp::FullPivLu,
1042            LinalgOp::Svd {
1043                derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
1044                gauge: SvdGauge::Raw,
1045            },
1046            LinalgOp::SvdVals {
1047                derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
1048            },
1049            LinalgOp::Qr {
1050                gauge: QrGauge::Raw,
1051            },
1052            LinalgOp::Eigh {
1053                derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
1054                gauge: EighGauge::Raw,
1055            },
1056            LinalgOp::EighVals {
1057                derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
1058            },
1059            LinalgOp::Eig {
1060                input_dtype: DType::F64,
1061            },
1062            LinalgOp::EigVals {
1063                input_dtype: DType::F64,
1064            },
1065        ];
1066
1067        for (case_index, kind) in cases.into_iter().enumerate() {
1068            let mut builder = GraphBuilder::<StdTensorOp>::new();
1069            let mut ctx = ShapeGuardContext::default();
1070            let primal = input_key(100 + case_index as u64);
1071            insert_meta(&mut ctx, primal.clone(), &[3]);
1072            let tangent = builder.add_input(TensorInputKey::User {
1073                id: 200 + case_index as u64,
1074            });
1075            let op = LinalgExtensionOp::new(kind);
1076            let outputs: Vec<_> = (0..op.output_count())
1077                .map(|offset| input_key(300 + case_index as u64 * 10 + offset as u64))
1078                .collect();
1079
1080            let result = LinalgAdRule
1081                .linearize(
1082                    &op,
1083                    &mut builder,
1084                    &[primal],
1085                    &outputs,
1086                    &[Some(tangent)],
1087                    &mut ctx,
1088                )
1089                .unwrap();
1090
1091            assert_eq!(result, vec![None; op.output_count()], "{kind:?}");
1092            assert!(
1093                builder.build().operations().is_empty(),
1094                "{kind:?} should not emit a malformed matrix AD graph"
1095            );
1096        }
1097    }
1098}
tenferro_linalg/ad.rs

tenferro_linalg/
ad.rs
