tenferro_linalg/
traced.rs

1use std::sync::Arc;
2
3use num_complex::{Complex32, Complex64};
4use tenferro_runtime::extension::apply;
5use tenferro_runtime::{CompareDir, DType, DotGeneralConfig, Error, Result, TracedTensor};
6
7use crate::extension::{
8    validate_derivative_eps, EighOptions, LinalgExtensionOp, LinalgOp, QrOptions, SvdOptions,
9};
10
11/// Linear algebra extension methods for [`TracedTensor`].
12pub trait TracedTensorLinalgExt {
13    fn svd(&self) -> Result<(TracedTensor, TracedTensor, TracedTensor)>;
14    fn svd_with_options(
15        &self,
16        options: SvdOptions,
17    ) -> Result<(TracedTensor, TracedTensor, TracedTensor)>;
18    fn qr(&self) -> Result<(TracedTensor, TracedTensor)>;
19    fn qr_with_options(&self, options: QrOptions) -> Result<(TracedTensor, TracedTensor)>;
20    fn eigh(&self) -> Result<(TracedTensor, TracedTensor)>;
21    fn eigh_with_options(&self, options: EighOptions) -> Result<(TracedTensor, TracedTensor)>;
22    fn cholesky(&self) -> Result<TracedTensor>;
23    fn lu(&self) -> Result<(TracedTensor, TracedTensor, TracedTensor, TracedTensor)>;
24    fn full_piv_lu(
25        &self,
26    ) -> Result<(
27        TracedTensor,
28        TracedTensor,
29        TracedTensor,
30        TracedTensor,
31        TracedTensor,
32    )>;
33    fn eig(&self) -> Result<(TracedTensor, TracedTensor)>;
34    fn solve(&self, b: &TracedTensor) -> Result<TracedTensor>;
35    fn full_piv_lu_solve(&self, b: &TracedTensor) -> Result<TracedTensor>;
36    fn triangular_solve(
37        &self,
38        b: &TracedTensor,
39        left_side: bool,
40        lower: bool,
41        transpose_a: bool,
42        unit_diagonal: bool,
43    ) -> Result<TracedTensor>;
44    fn slogdet(&self) -> Result<(TracedTensor, TracedTensor)>;
45    fn det(&self) -> Result<TracedTensor>;
46    fn inv(&self) -> Result<TracedTensor>;
47    fn eigvalsh(&self) -> Result<TracedTensor>;
48    fn eigvals(&self) -> Result<TracedTensor>;
49    fn pinv(&self) -> Result<TracedTensor>;
50    fn pinv_with_rtol(&self, rtol: f64) -> Result<TracedTensor>;
51    fn norm(&self, ord: Option<f64>, dim: Option<&[usize]>, keepdim: bool) -> Result<TracedTensor>;
52}
53
54impl TracedTensorLinalgExt for TracedTensor {
55    fn svd(&self) -> Result<(TracedTensor, TracedTensor, TracedTensor)> {
56        svd(self)
57    }
58
59    fn svd_with_options(
60        &self,
61        options: SvdOptions,
62    ) -> Result<(TracedTensor, TracedTensor, TracedTensor)> {
63        svd_with_options(self, options)
64    }
65
66    fn qr(&self) -> Result<(TracedTensor, TracedTensor)> {
67        qr(self)
68    }
69
70    fn qr_with_options(&self, options: QrOptions) -> Result<(TracedTensor, TracedTensor)> {
71        qr_with_options(self, options)
72    }
73
74    fn eigh(&self) -> Result<(TracedTensor, TracedTensor)> {
75        eigh(self)
76    }
77
78    fn eigh_with_options(&self, options: EighOptions) -> Result<(TracedTensor, TracedTensor)> {
79        eigh_with_options(self, options)
80    }
81
82    fn cholesky(&self) -> Result<TracedTensor> {
83        cholesky(self)
84    }
85
86    fn lu(&self) -> Result<(TracedTensor, TracedTensor, TracedTensor, TracedTensor)> {
87        lu(self)
88    }
89
90    fn full_piv_lu(
91        &self,
92    ) -> Result<(
93        TracedTensor,
94        TracedTensor,
95        TracedTensor,
96        TracedTensor,
97        TracedTensor,
98    )> {
99        full_piv_lu(self)
100    }
101
102    fn eig(&self) -> Result<(TracedTensor, TracedTensor)> {
103        eig(self)
104    }
105
106    fn solve(&self, b: &TracedTensor) -> Result<TracedTensor> {
107        solve(self, b)
108    }
109
110    fn full_piv_lu_solve(&self, b: &TracedTensor) -> Result<TracedTensor> {
111        full_piv_lu_solve(self, b)
112    }
113
114    fn triangular_solve(
115        &self,
116        b: &TracedTensor,
117        left_side: bool,
118        lower: bool,
119        transpose_a: bool,
120        unit_diagonal: bool,
121    ) -> Result<TracedTensor> {
122        triangular_solve(self, b, left_side, lower, transpose_a, unit_diagonal)
123    }
124
125    fn slogdet(&self) -> Result<(TracedTensor, TracedTensor)> {
126        slogdet(self)
127    }
128
129    fn det(&self) -> Result<TracedTensor> {
130        det(self)
131    }
132
133    fn inv(&self) -> Result<TracedTensor> {
134        inv(self)
135    }
136
137    fn eigvalsh(&self) -> Result<TracedTensor> {
138        eigvalsh(self)
139    }
140
141    fn eigvals(&self) -> Result<TracedTensor> {
142        eigvals(self)
143    }
144
145    fn pinv(&self) -> Result<TracedTensor> {
146        pinv(self)
147    }
148
149    fn pinv_with_rtol(&self, rtol: f64) -> Result<TracedTensor> {
150        pinv_with_rtol(self, rtol)
151    }
152
153    fn norm(&self, ord: Option<f64>, dim: Option<&[usize]>, keepdim: bool) -> Result<TracedTensor> {
154        norm(self, ord, dim, keepdim)
155    }
156}
157
158/// Build a traced singular value decomposition op using default options.
159///
160/// # Examples
161///
162/// ```
163/// use tenferro_linalg::TracedTensorLinalgExt;
164/// use tenferro_runtime::TracedTensor;
165///
166/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap();
167/// let (u, s, vt) = a.svd().unwrap();
168/// assert_eq!(u.rank, 2);
169/// assert_eq!(s.rank, 1);
170/// assert_eq!(vt.rank, 2);
171/// ```
172pub fn svd(a: &TracedTensor) -> Result<(TracedTensor, TracedTensor, TracedTensor)> {
173    svd_with_options(a, SvdOptions::default())
174}
175
176/// Build a traced singular value decomposition op with explicit options.
177///
178/// `derivative_eps` regularizes decomposition derivative formulas. It is not a
179/// backend SVD solver tolerance.
180///
181/// # Examples
182///
183/// ```
184/// use tenferro_linalg::{SvdGauge, SvdOptions, TracedTensorLinalgExt};
185/// use tenferro_runtime::TracedTensor;
186///
187/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap();
188/// let options = SvdOptions::default()
189///     .gauge(SvdGauge::CanonicalPivot)
190///     .derivative_eps(1e-10);
191/// let (_u, s, _vt) = a.svd_with_options(options).unwrap();
192/// assert_eq!(s.rank, 1);
193/// ```
194pub fn svd_with_options(
195    a: &TracedTensor,
196    options: SvdOptions,
197) -> Result<(TracedTensor, TracedTensor, TracedTensor)> {
198    validate_derivative_eps("svd_with_options", options.derivative_eps)?;
199    three_outputs(
200        apply(
201            Arc::new(LinalgExtensionOp::new(LinalgOp::Svd {
202                derivative_eps: options.derivative_eps,
203                gauge: options.gauge,
204            })),
205            &[a],
206        )?,
207        "svd",
208    )
209}
210
211/// Build a traced QR decomposition op.
212///
213/// # Examples
214///
215/// ```
216/// use tenferro_linalg::TracedTensorLinalgExt;
217/// use tenferro_runtime::TracedTensor;
218///
219/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap();
220/// let (q, r) = a.qr().unwrap();
221/// assert_eq!(q.rank, 2);
222/// assert_eq!(r.rank, 2);
223/// ```
224pub fn qr(a: &TracedTensor) -> Result<(TracedTensor, TracedTensor)> {
225    qr_with_options(a, QrOptions::default())
226}
227
228/// Build a traced QR decomposition op with explicit options.
229///
230/// `gauge` controls optional sign or phase post-processing.
231///
232/// # Examples
233///
234/// ```
235/// use tenferro_linalg::{QrGauge, QrOptions, TracedTensorLinalgExt};
236/// use tenferro_runtime::TracedTensor;
237///
238/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap();
239/// let (q, r) = a.qr_with_options(QrOptions::default().gauge(QrGauge::PositiveDiagonal)).unwrap();
240/// assert_eq!(q.rank, 2);
241/// assert_eq!(r.rank, 2);
242/// ```
243pub fn qr_with_options(
244    a: &TracedTensor,
245    options: QrOptions,
246) -> Result<(TracedTensor, TracedTensor)> {
247    two_outputs(
248        apply(
249            Arc::new(LinalgExtensionOp::new(LinalgOp::Qr {
250                gauge: options.gauge,
251            })),
252            &[a],
253        )?,
254        "qr",
255    )
256}
257
258/// Build a traced Hermitian eigenvalue decomposition op using default options.
259///
260/// # Examples
261///
262/// ```
263/// use tenferro_linalg::TracedTensorLinalgExt;
264/// use tenferro_runtime::TracedTensor;
265///
266/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 3.0]).unwrap();
267/// let (values, vectors) = a.eigh().unwrap();
268/// assert_eq!(values.rank, 1);
269/// assert_eq!(vectors.rank, 2);
270/// ```
271pub fn eigh(a: &TracedTensor) -> Result<(TracedTensor, TracedTensor)> {
272    eigh_with_options(a, EighOptions::default())
273}
274
275/// Build a traced Hermitian eigenvalue decomposition op with explicit options.
276///
277/// `derivative_eps` regularizes derivative formulas for repeated or nearly
278/// repeated eigenvalues. It is not a backend eigensolver tolerance.
279///
280/// # Examples
281///
282/// ```
283/// use tenferro_linalg::{EighGauge, EighOptions, TracedTensorLinalgExt};
284/// use tenferro_runtime::TracedTensor;
285///
286/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 3.0]).unwrap();
287/// let (values, _vectors) = a
288///     .eigh_with_options(
289///         EighOptions::default()
290///             .gauge(EighGauge::CanonicalPivot)
291///             .derivative_eps(1e-10),
292///     )
293///     .unwrap();
294/// assert_eq!(values.rank, 1);
295/// ```
296pub fn eigh_with_options(
297    a: &TracedTensor,
298    options: EighOptions,
299) -> Result<(TracedTensor, TracedTensor)> {
300    validate_derivative_eps("eigh_with_options", options.derivative_eps)?;
301    two_outputs(
302        apply(
303            Arc::new(LinalgExtensionOp::new(LinalgOp::Eigh {
304                derivative_eps: options.derivative_eps,
305                gauge: options.gauge,
306            })),
307            &[a],
308        )?,
309        "eigh",
310    )
311}
312
313/// Build a traced Cholesky decomposition op.
314///
315/// # Examples
316///
317/// ```
318/// use tenferro_linalg::TracedTensorLinalgExt;
319/// use tenferro_runtime::TracedTensor;
320///
321/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![4.0_f64, 2.0, 2.0, 3.0]).unwrap();
322/// let factor = a.cholesky().unwrap();
323/// assert_eq!(factor.rank, 2);
324/// ```
325pub fn cholesky(a: &TracedTensor) -> Result<TracedTensor> {
326    one_output(
327        apply(Arc::new(LinalgExtensionOp::new(LinalgOp::Cholesky)), &[a])?,
328        "cholesky",
329    )
330}
331
332/// Build a traced LU decomposition op.
333///
334/// # Examples
335///
336/// ```
337/// use tenferro_linalg::TracedTensorLinalgExt;
338/// use tenferro_runtime::TracedTensor;
339///
340/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 3.0, 2.0, 4.0]).unwrap();
341/// let (p, l, u, parity) = a.lu().unwrap();
342/// assert_eq!(p.rank, 2);
343/// assert_eq!(l.rank, 2);
344/// assert_eq!(u.rank, 2);
345/// assert_eq!(parity.rank, 0);
346/// ```
347pub fn lu(a: &TracedTensor) -> Result<(TracedTensor, TracedTensor, TracedTensor, TracedTensor)> {
348    four_outputs(
349        apply(Arc::new(LinalgExtensionOp::new(LinalgOp::Lu)), &[a])?,
350        "lu",
351    )
352}
353
354/// Build a traced full-pivot LU decomposition op.
355///
356/// Returns `(P, L, U, Q, parity)` with reconstruction convention
357/// `A = P^T * L * U * Q`, equivalently `P * A * Q^T = L * U`. `parity` is a
358/// scalar real tensor containing `+1` or `-1`: `F32` for `F32`/`C32` inputs and
359/// `F64` for `F64`/`C64` inputs.
360///
361/// # Examples
362///
363/// ```
364/// use tenferro_linalg::TracedTensorLinalgExt;
365/// use tenferro_runtime::TracedTensor;
366///
367/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 3.0, 2.0, 4.0]).unwrap();
368/// let (p, l, u, q, parity) = a.full_piv_lu().unwrap();
369/// assert_eq!(p.rank, 2);
370/// assert_eq!(l.rank, 2);
371/// assert_eq!(u.rank, 2);
372/// assert_eq!(q.rank, 2);
373/// assert_eq!(parity.rank, 0);
374/// ```
375pub fn full_piv_lu(
376    a: &TracedTensor,
377) -> Result<(
378    TracedTensor,
379    TracedTensor,
380    TracedTensor,
381    TracedTensor,
382    TracedTensor,
383)> {
384    five_outputs(
385        apply(Arc::new(LinalgExtensionOp::new(LinalgOp::FullPivLu)), &[a])?,
386        "full_piv_lu",
387    )
388}
389
390/// Build a traced general eigendecomposition op.
391///
392/// # Examples
393///
394/// ```
395/// use tenferro_linalg::TracedTensorLinalgExt;
396/// use tenferro_runtime::TracedTensor;
397///
398/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0]).unwrap();
399/// let (values, vectors) = a.eig().unwrap();
400/// assert_eq!(values.rank, 1);
401/// assert_eq!(vectors.rank, 2);
402/// ```
403pub fn eig(a: &TracedTensor) -> Result<(TracedTensor, TracedTensor)> {
404    two_outputs(
405        apply(
406            Arc::new(LinalgExtensionOp::new(LinalgOp::Eig {
407                input_dtype: a.dtype,
408            })),
409            &[a],
410        )?,
411        "eig",
412    )
413}
414
415/// Build a traced linear solve op.
416///
417/// # Examples
418///
419/// ```
420/// use tenferro_linalg::TracedTensorLinalgExt;
421/// use tenferro_runtime::TracedTensor;
422///
423/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 3.0]).unwrap();
424/// let b = TracedTensor::from_vec_col_major(vec![2, 1], vec![4.0_f64, 9.0]).unwrap();
425/// let x = a.solve(&b).unwrap();
426/// assert_eq!(x.rank, 2);
427/// ```
428pub fn solve(a: &TracedTensor, b: &TracedTensor) -> Result<TracedTensor> {
429    let mut factor_outputs =
430        apply(Arc::new(LinalgExtensionOp::new(LinalgOp::LuFactor)), &[a])?.into_iter();
431    let (packed_lu, pivots) = match (
432        factor_outputs.next(),
433        factor_outputs.next(),
434        factor_outputs.next(),
435        factor_outputs.next(),
436    ) {
437        (Some(packed_lu), Some(pivots), Some(_parity), None) => (packed_lu, pivots),
438        _ => return Err(unexpected_output_count("lu_factor", 3)),
439    };
440    one_output(
441        apply(
442            Arc::new(LinalgExtensionOp::new(LinalgOp::LuSolvePrepared {
443                transpose_a: false,
444                conjugate_a: false,
445            })),
446            &[a, &packed_lu, &pivots, b],
447        )?,
448        "solve",
449    )
450}
451
452/// Build a traced full-pivot LU solve op.
453///
454/// # Examples
455///
456/// ```
457/// use tenferro_linalg::TracedTensorLinalgExt;
458/// use tenferro_runtime::TracedTensor;
459///
460/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 3.0]).unwrap();
461/// let b = TracedTensor::from_vec_col_major(vec![2, 1], vec![4.0_f64, 9.0]).unwrap();
462/// let x = a.full_piv_lu_solve(&b).unwrap();
463/// assert_eq!(x.rank, 2);
464/// ```
465pub fn full_piv_lu_solve(a: &TracedTensor, b: &TracedTensor) -> Result<TracedTensor> {
466    one_output(
467        apply(
468            Arc::new(LinalgExtensionOp::new(LinalgOp::FullPivLuSolve {
469                transpose_a: false,
470            })),
471            &[a, b],
472        )?,
473        "full_piv_lu_solve",
474    )
475}
476
477/// Build a traced triangular solve op.
478///
479/// # Examples
480///
481/// ```
482/// use tenferro_linalg::TracedTensorLinalgExt;
483/// use tenferro_runtime::TracedTensor;
484///
485/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 1.0, 3.0]).unwrap();
486/// let b = TracedTensor::from_vec_col_major(vec![2, 1], vec![4.0_f64, 9.0]).unwrap();
487/// let x = a.triangular_solve(&b, true, true, false, false).unwrap();
488/// assert_eq!(x.rank, 2);
489/// ```
490pub fn triangular_solve(
491    a: &TracedTensor,
492    b: &TracedTensor,
493    left_side: bool,
494    lower: bool,
495    transpose_a: bool,
496    unit_diagonal: bool,
497) -> Result<TracedTensor> {
498    one_output(
499        apply(
500            Arc::new(LinalgExtensionOp::new(LinalgOp::TriangularSolve {
501                left_side,
502                lower,
503                transpose_a,
504                unit_diagonal,
505            })),
506            &[a, b],
507        )?,
508        "triangular_solve",
509    )
510}
511
512/// Build traced sign and log-absolute-determinant ops.
513///
514/// # Examples
515///
516/// ```
517/// use tenferro_linalg::TracedTensorLinalgExt;
518/// use tenferro_runtime::TracedTensor;
519///
520/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 3.0]).unwrap();
521/// let (sign, logabsdet) = a.slogdet().unwrap();
522/// assert_eq!(sign.rank, 0);
523/// assert_eq!(logabsdet.rank, 0);
524/// ```
525pub fn slogdet(a: &TracedTensor) -> Result<(TracedTensor, TracedTensor)> {
526    let mut factor_outputs =
527        apply(Arc::new(LinalgExtensionOp::new(LinalgOp::LuFactor)), &[a])?.into_iter();
528    let (packed_lu, parity) = match (
529        factor_outputs.next(),
530        factor_outputs.next(),
531        factor_outputs.next(),
532        factor_outputs.next(),
533    ) {
534        (Some(packed_lu), Some(_pivots), Some(parity), None) => (packed_lu, parity),
535        _ => return Err(unexpected_output_count("lu_factor", 3)),
536    };
537    let diag_u = packed_lu.extract_diag(0, 1)?;
538    let sign_u = diag_u.sign()?.reduce_prod(&[0])?;
539    let sign = (&parity * &sign_u)?;
540    let logabsdet = diag_u.abs()?.log()?.reduce_sum(&[0])?;
541    Ok((sign, logabsdet))
542}
543
544/// Build a traced determinant op.
545///
546/// # Examples
547///
548/// ```
549/// use tenferro_linalg::TracedTensorLinalgExt;
550/// use tenferro_runtime::TracedTensor;
551///
552/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 3.0]).unwrap();
553/// let determinant = a.det().unwrap();
554/// assert_eq!(determinant.rank, 0);
555/// ```
556pub fn det(a: &TracedTensor) -> Result<TracedTensor> {
557    let (sign, logabsdet) = slogdet(a)?;
558    &sign * &logabsdet.exp()?
559}
560
561/// Build a traced matrix inverse op.
562///
563/// # Examples
564///
565/// ```
566/// use tenferro_linalg::TracedTensorLinalgExt;
567/// use tenferro_runtime::TracedTensor;
568///
569/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 3.0]).unwrap();
570/// let inverse = a.inv().unwrap();
571/// assert_eq!(inverse.rank, 2);
572/// ```
573pub fn inv(a: &TracedTensor) -> Result<TracedTensor> {
574    ensure_min_rank("inv", a.rank, 2)?;
575    let shape = require_concrete_shape("inv", a)?;
576    let eye = eye_like(a, shape[0])?;
577    solve(a, &eye)
578}
579
580/// Build a traced Hermitian eigenvalue-only op.
581///
582/// # Examples
583///
584/// ```
585/// use tenferro_linalg::TracedTensorLinalgExt;
586/// use tenferro_runtime::TracedTensor;
587///
588/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 3.0]).unwrap();
589/// let values = a.eigvalsh().unwrap();
590/// assert_eq!(values.rank, 1);
591/// ```
592pub fn eigvalsh(a: &TracedTensor) -> Result<TracedTensor> {
593    eigh_values(a)
594}
595
596/// Build a traced general eigenvalue-only op.
597///
598/// # Examples
599///
600/// ```
601/// use tenferro_linalg::TracedTensorLinalgExt;
602/// use tenferro_runtime::TracedTensor;
603///
604/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0]).unwrap();
605/// let values = a.eigvals().unwrap();
606/// assert_eq!(values.rank, 1);
607/// ```
608pub fn eigvals(a: &TracedTensor) -> Result<TracedTensor> {
609    eig_values(a)
610}
611
612/// Build a traced Moore-Penrose pseudoinverse op.
613///
614/// Floating-point and complex inputs are supported. Integer and boolean inputs
615/// return an unsupported-dtype error.
616///
617/// # Examples
618///
619/// ```
620/// use tenferro_linalg::TracedTensorLinalgExt;
621/// use tenferro_runtime::TracedTensor;
622///
623/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0]).unwrap();
624/// let inverse = a.pinv().unwrap();
625/// assert_eq!(inverse.rank, 2);
626/// ```
627pub fn pinv(a: &TracedTensor) -> Result<TracedTensor> {
628    ensure_float_or_complex("pinv", a.dtype)?;
629    let shape = require_concrete_shape("pinv", a)?;
630    let max_dim = match (shape.first(), shape.get(1)) {
631        (Some(&m), Some(&n)) => m.max(n),
632        (Some(&m), None) => m,
633        _ => 0,
634    };
635    pinv_with_rtol(a, default_pinv_rtol(a.dtype, max_dim))
636}
637
638/// Build a traced Moore-Penrose pseudoinverse op with an explicit relative tolerance.
639///
640/// Floating-point and complex inputs are supported. Integer and boolean inputs
641/// return an unsupported-dtype error.
642///
643/// # Examples
644///
645/// ```
646/// use tenferro_linalg::TracedTensorLinalgExt;
647/// use tenferro_runtime::TracedTensor;
648///
649/// let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0]).unwrap();
650/// let inverse = a.pinv_with_rtol(1e-12).unwrap();
651/// assert_eq!(inverse.rank, 2);
652/// ```
653pub fn pinv_with_rtol(a: &TracedTensor, rtol: f64) -> Result<TracedTensor> {
654    ensure_float_or_complex("pinv_with_rtol", a.dtype)?;
655    require_concrete_shape("pinv_with_rtol", a)?;
656    let (u, s, vt) = svd(a)?;
657    let abs_s = s.abs()?;
658    let s_max = abs_s.reduce_max(&[0])?;
659    let s_max_shape = s_max.concrete_shape()?;
660    let threshold_scalar = broadcast_scalar(scalar_real(s.dtype, rtol.max(0.0))?, &s_max_shape)?;
661    let threshold = (&s_max * &threshold_scalar)?;
662    let s_shape = s.concrete_shape()?;
663    let threshold = broadcast_batch_scalar_to_leading_axis(&threshold, &s_shape)?;
664    let mask = abs_s.compare(&threshold, CompareDir::Gt)?;
665    let mask = mask.convert(s.dtype)?;
666    let ones = ones_like(&s)?;
667    let neg_mask = (-&mask)?;
668    let denom = (&s + &(&ones + &neg_mask)?)?;
669    let s_inv = (&mask / &denom)?;
670
671    let v = vt.conj()?.transpose(&matrix_transpose_perm(vt.rank))?;
672    let uh = u.conj()?.transpose(&matrix_transpose_perm(u.rank))?;
673    let vs = scale_matrix_columns(&v, &s_inv)?;
674    matmul_preserve_trailing_batch(&vs, &uh)
675}
676
677/// Build a traced vector, matrix, or tensor norm op.
678///
679/// Floating-point and complex inputs are supported. Integer and boolean inputs
680/// return an unsupported-dtype error.
681///
682/// # Examples
683///
684/// ```
685/// use tenferro_linalg::TracedTensorLinalgExt;
686/// use tenferro_runtime::TracedTensor;
687///
688/// let x = TracedTensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap();
689/// let length = x.norm(Some(2.0), Some(&[0]), false).unwrap();
690/// assert_eq!(length.rank, 0);
691/// ```
692pub fn norm(
693    a: &TracedTensor,
694    ord: Option<f64>,
695    dim: Option<&[usize]>,
696    keepdim: bool,
697) -> Result<TracedTensor> {
698    ensure_float_or_complex("norm", a.dtype)?;
699    let shape = require_concrete_shape("norm", a)?;
700    let axes = dim.map_or_else(|| (0..a.rank).collect::<Vec<_>>(), |dims| dims.to_vec());
701    if axes.is_empty() {
702        return Ok(a.clone());
703    }
704    validate_axes("norm", a.rank, &axes)?;
705
706    let out = match axes.len() {
707        1 => vector_norm(a, axes[0], ord)?,
708        2 => matrix_norm(a, &axes, ord)?,
709        _ => {
710            let abs = a.abs()?;
711            match ord {
712                None => frobenius_norm(&abs, &axes)?,
713                Some(p) if p == f64::INFINITY => abs.reduce_max(&axes)?,
714                Some(p) if p == f64::NEG_INFINITY => abs.reduce_min(&axes)?,
715                Some(0.0) => count_nonzero(&abs, &axes)?,
716                Some(p) => p_norm(&abs, &axes, p)?,
717            }
718        }
719    };
720    restore_keepdim(out, &shape, &axes, keepdim)
721}
722
723fn unexpected_output_count(name: &str, expected: usize) -> Error {
724    Error::Internal(format!("{name} must produce exactly {expected} outputs"))
725}
726
727fn one_output(outputs: Vec<TracedTensor>, name: &str) -> Result<TracedTensor> {
728    let mut outputs = outputs.into_iter();
729    match (outputs.next(), outputs.next()) {
730        (Some(output), None) => Ok(output),
731        _ => Err(unexpected_output_count(name, 1)),
732    }
733}
734
735fn two_outputs(outputs: Vec<TracedTensor>, name: &str) -> Result<(TracedTensor, TracedTensor)> {
736    let mut outputs = outputs.into_iter();
737    match (outputs.next(), outputs.next(), outputs.next()) {
738        (Some(lhs), Some(rhs), None) => Ok((lhs, rhs)),
739        _ => Err(unexpected_output_count(name, 2)),
740    }
741}
742
743fn three_outputs(
744    outputs: Vec<TracedTensor>,
745    name: &str,
746) -> Result<(TracedTensor, TracedTensor, TracedTensor)> {
747    let mut outputs = outputs.into_iter();
748    match (
749        outputs.next(),
750        outputs.next(),
751        outputs.next(),
752        outputs.next(),
753    ) {
754        (Some(first), Some(second), Some(third), None) => Ok((first, second, third)),
755        _ => Err(unexpected_output_count(name, 3)),
756    }
757}
758
759fn four_outputs(
760    outputs: Vec<TracedTensor>,
761    name: &str,
762) -> Result<(TracedTensor, TracedTensor, TracedTensor, TracedTensor)> {
763    let mut outputs = outputs.into_iter();
764    match (
765        outputs.next(),
766        outputs.next(),
767        outputs.next(),
768        outputs.next(),
769        outputs.next(),
770    ) {
771        (Some(first), Some(second), Some(third), Some(fourth), None) => {
772            Ok((first, second, third, fourth))
773        }
774        _ => Err(unexpected_output_count(name, 4)),
775    }
776}
777
778fn five_outputs(
779    outputs: Vec<TracedTensor>,
780    name: &str,
781) -> Result<(
782    TracedTensor,
783    TracedTensor,
784    TracedTensor,
785    TracedTensor,
786    TracedTensor,
787)> {
788    let mut outputs = outputs.into_iter();
789    match (
790        outputs.next(),
791        outputs.next(),
792        outputs.next(),
793        outputs.next(),
794        outputs.next(),
795        outputs.next(),
796    ) {
797        (Some(first), Some(second), Some(third), Some(fourth), Some(fifth), None) => {
798            Ok((first, second, third, fourth, fifth))
799        }
800        _ => Err(unexpected_output_count(name, 5)),
801    }
802}
803
804fn scalar_real(dtype: DType, value: f64) -> Result<TracedTensor> {
805    match dtype {
806        DType::F64 => TracedTensor::from_vec_col_major(vec![], vec![value]),
807        DType::F32 => TracedTensor::from_vec_col_major(vec![], vec![value as f32]),
808        DType::I32 => TracedTensor::from_vec_col_major(vec![], vec![value.round() as i32]),
809        DType::I64 => TracedTensor::from_vec_col_major(vec![], vec![value.round() as i64]),
810        DType::Bool => TracedTensor::from_vec_col_major(vec![], vec![value != 0.0]),
811        DType::C64 => TracedTensor::from_vec_col_major(vec![], vec![Complex64::new(value, 0.0)]),
812        DType::C32 => {
813            TracedTensor::from_vec_col_major(vec![], vec![Complex32::new(value as f32, 0.0)])
814        }
815    }
816}
817
818fn ensure_float_or_complex(op: &'static str, dtype: DType) -> Result<()> {
819    match dtype {
820        DType::F32 | DType::F64 | DType::C32 | DType::C64 => Ok(()),
821        DType::I32 | DType::I64 | DType::Bool => Err(Error::TensorRuntime(
822            tenferro_tensor::Error::backend_failure(op, format!("unsupported dtype {dtype:?}")),
823        )),
824    }
825}
826
827fn ensure_min_rank(op: &'static str, actual: usize, expected: usize) -> Result<()> {
828    if actual < expected {
829        return Err(Error::TensorRuntime(tenferro_tensor::Error::RankMismatch {
830            op,
831            expected,
832            actual,
833        }));
834    }
835    Ok(())
836}
837
838fn validate_axes(op: &'static str, rank: usize, axes: &[usize]) -> Result<()> {
839    for &axis in axes {
840        if axis >= rank {
841            return Err(Error::TensorRuntime(
842                tenferro_tensor::Error::AxisOutOfBounds { op, axis, rank },
843            ));
844        }
845    }
846    Ok(())
847}
848
849fn require_concrete_shape(op: &'static str, input: &TracedTensor) -> Result<Vec<usize>> {
850    input.try_concrete_shape().ok_or_else(|| {
851        Error::TensorRuntime(tenferro_tensor::Error::backend_failure(
852            op,
853            "symbolic shape is not supported by this traced linalg helper",
854        ))
855    })
856}
857
858fn zero_scalar(dtype: DType) -> Result<TracedTensor> {
859    scalar_real(dtype, 0.0)
860}
861
862fn one_scalar(dtype: DType) -> Result<TracedTensor> {
863    scalar_real(dtype, 1.0)
864}
865
866fn ones_like(input: &TracedTensor) -> Result<TracedTensor> {
867    let shape = input.concrete_shape()?;
868    broadcast_scalar(one_scalar(input.dtype)?, &shape)
869}
870
871fn eye_like(anchor: &TracedTensor, size: usize) -> Result<TracedTensor> {
872    let mut vector_shape = vec![size];
873    let anchor_shape = anchor.concrete_shape()?;
874    vector_shape.extend_from_slice(&anchor_shape[2..]);
875    let diagonal = broadcast_scalar(one_scalar(anchor.dtype)?, &vector_shape)?;
876    diagonal.embed_diag(0, 1)
877}
878
879fn broadcast_scalar(input: TracedTensor, shape: &[usize]) -> Result<TracedTensor> {
880    let input_shape = input.concrete_shape()?;
881    if input_shape == shape {
882        return Ok(input);
883    }
884    input.broadcast_in_dim(shape, &[])
885}
886
887fn broadcast_batch_scalar_to_leading_axis(
888    input: &TracedTensor,
889    shape: &[usize],
890) -> Result<TracedTensor> {
891    let input_shape = input.concrete_shape()?;
892    if input_shape == shape {
893        return Ok(input.clone());
894    }
895    let dims: Vec<usize> = (1..shape.len()).collect();
896    input.broadcast_in_dim(shape, &dims)
897}
898
899fn matmul_preserve_trailing_batch(lhs: &TracedTensor, rhs: &TracedTensor) -> Result<TracedTensor> {
900    let rank = lhs.rank;
901    let batch_dims: Vec<usize> = (2..rank).collect();
902    lhs.dot_general(
903        rhs,
904        DotGeneralConfig {
905            lhs_contracting_dims: vec![1],
906            rhs_contracting_dims: vec![0],
907            lhs_batch_dims: batch_dims.clone(),
908            rhs_batch_dims: batch_dims,
909        },
910    )
911}
912
913fn matrix_transpose_perm(rank: usize) -> Vec<usize> {
914    let mut perm: Vec<usize> = (0..rank).collect();
915    perm.swap(0, 1);
916    perm
917}
918
919fn frobenius_norm(abs: &TracedTensor, axes: &[usize]) -> Result<TracedTensor> {
920    let squared = abs.pow(&scalar_real(abs.dtype, 2.0)?)?;
921    squared.reduce_sum(axes)?.sqrt()
922}
923
924fn p_norm(abs: &TracedTensor, axes: &[usize], p: f64) -> Result<TracedTensor> {
925    if !p.is_finite() || p == 0.0 {
926        return Err(Error::InvalidGraphBuild {
927            op: "norm",
928            message: format!("p-norm order must be finite and nonzero, got {p}"),
929        });
930    }
931    let power = abs.pow(&scalar_real(abs.dtype, p)?)?;
932    let inv_p = scalar_real(abs.dtype, 1.0 / p)?;
933    power.reduce_sum(axes)?.pow(&inv_p)
934}
935
936fn default_pinv_rtol(dtype: DType, max_dim: usize) -> f64 {
937    let eps = match dtype {
938        DType::F32 | DType::C32 => f32::EPSILON as f64,
939        DType::F64 | DType::C64 => f64::EPSILON,
940        DType::I32 | DType::I64 | DType::Bool => 0.0,
941    };
942    eps * max_dim as f64
943}
944
945fn vector_norm(a: &TracedTensor, axis: usize, ord: Option<f64>) -> Result<TracedTensor> {
946    let abs = a.abs()?;
947    match ord {
948        None => frobenius_norm(&abs, &[axis]),
949        Some(0.0) => count_nonzero(&abs, &[axis]),
950        Some(p) if p == f64::INFINITY => abs.reduce_max(&[axis]),
951        Some(p) if p == f64::NEG_INFINITY => abs.reduce_min(&[axis]),
952        Some(p) => p_norm(&abs, &[axis], p),
953    }
954}
955
956fn matrix_norm(a: &TracedTensor, axes: &[usize], ord: Option<f64>) -> Result<TracedTensor> {
957    let matrix = move_axes_to_front(a, axes)?;
958    let abs = matrix.abs()?;
959    match ord {
960        None => frobenius_norm(&abs, &[0, 1]),
961        Some(p) if p == f64::INFINITY => matrix_row_sum_norm(&abs, true),
962        Some(p) if p == f64::NEG_INFINITY => matrix_row_sum_norm(&abs, false),
963        Some(1.0) => matrix_col_sum_norm(&abs, true),
964        Some(-1.0) => matrix_col_sum_norm(&abs, false),
965        Some(2.0) => {
966            let singular_values = svd_values(&matrix)?.abs()?;
967            singular_values.reduce_max(&[0])
968        }
969        Some(-2.0) => {
970            let singular_values = svd_values(&matrix)?.abs()?;
971            singular_values.reduce_min(&[0])
972        }
973        Some(0.0) => count_nonzero(&abs, &[0, 1]),
974        Some(p) => p_norm(&abs, &[0, 1], p),
975    }
976}
977
978fn svd_values(a: &TracedTensor) -> Result<TracedTensor> {
979    let (_u, s, _vt) = three_outputs(
980        apply(
981            Arc::new(LinalgExtensionOp::new(LinalgOp::Svd {
982                derivative_eps: SvdOptions::default().derivative_eps,
983                gauge: SvdOptions::default().gauge,
984            })),
985            &[a],
986        )?,
987        "svd_values",
988    )?;
989    Ok(s)
990}
991
992fn eigh_values(a: &TracedTensor) -> Result<TracedTensor> {
993    let (values, _vectors) = two_outputs(
994        apply(
995            Arc::new(LinalgExtensionOp::new(LinalgOp::Eigh {
996                derivative_eps: EighOptions::default().derivative_eps,
997                gauge: EighOptions::default().gauge,
998            })),
999            &[a],
1000        )?,
1001        "eigh_values",
1002    )?;
1003    Ok(values)
1004}
1005
1006fn eig_values(a: &TracedTensor) -> Result<TracedTensor> {
1007    let (values, _vectors) = two_outputs(
1008        apply(
1009            Arc::new(LinalgExtensionOp::new(LinalgOp::Eig {
1010                input_dtype: a.dtype,
1011            })),
1012            &[a],
1013        )?,
1014        "eig_values",
1015    )?;
1016    Ok(values)
1017}
1018
1019fn scale_matrix_columns(matrix: &TracedTensor, scale: &TracedTensor) -> Result<TracedTensor> {
1020    let matrix_shape = matrix.concrete_shape()?;
1021    let scale_shape_input = scale.concrete_shape()?;
1022    let mut scale_shape = vec![1, scale_shape_input[0]];
1023    scale_shape.extend_from_slice(&matrix_shape[2..]);
1024    let dims: Vec<usize> = (0..matrix_shape.len()).collect();
1025    let scale = scale
1026        .reshape(&scale_shape)?
1027        .broadcast_in_dim(&matrix_shape, &dims)?;
1028    matrix * &scale
1029}
1030
1031fn count_nonzero(abs: &TracedTensor, axes: &[usize]) -> Result<TracedTensor> {
1032    let mask = abs.compare(&zero_scalar(abs.dtype)?, CompareDir::Gt)?;
1033    mask.convert(abs.dtype)?.reduce_sum(axes)
1034}
1035
1036fn matrix_row_sum_norm(abs: &TracedTensor, take_max: bool) -> Result<TracedTensor> {
1037    let row_sums = abs.reduce_sum(&[1])?;
1038    if take_max {
1039        row_sums.reduce_max(&[0])
1040    } else {
1041        row_sums.reduce_min(&[0])
1042    }
1043}
1044
1045fn matrix_col_sum_norm(abs: &TracedTensor, take_max: bool) -> Result<TracedTensor> {
1046    let col_sums = abs.reduce_sum(&[0])?;
1047    if take_max {
1048        col_sums.reduce_max(&[0])
1049    } else {
1050        col_sums.reduce_min(&[0])
1051    }
1052}
1053
1054fn move_axes_to_front(tensor: &TracedTensor, axes: &[usize]) -> Result<TracedTensor> {
1055    if axes.iter().enumerate().all(|(index, &axis)| index == axis) {
1056        return Ok(tensor.clone());
1057    }
1058
1059    let mut selected = vec![false; tensor.rank];
1060    for &axis in axes {
1061        selected[axis] = true;
1062    }
1063
1064    let mut perm = Vec::with_capacity(tensor.rank);
1065    perm.extend_from_slice(axes);
1066    for (axis, is_selected) in selected.iter().enumerate().take(tensor.rank) {
1067        if !*is_selected {
1068            perm.push(axis);
1069        }
1070    }
1071    tensor.transpose(&perm)
1072}
1073
1074fn restore_keepdim(
1075    reduced: TracedTensor,
1076    original_shape: &[usize],
1077    axes: &[usize],
1078    keepdim: bool,
1079) -> Result<TracedTensor> {
1080    if !keepdim {
1081        return Ok(reduced);
1082    }
1083    let mut kept_shape = original_shape.to_vec();
1084    for &axis in axes {
1085        kept_shape[axis] = 1;
1086    }
1087    reduced.reshape(&kept_shape)
1088}
1089
1090#[cfg(test)]
1091mod tests {
1092    use super::p_norm;
1093    use tenferro_runtime::TracedTensor;
1094
1095    #[test]
1096    fn p_norm_rejects_zero_and_non_finite_orders() {
1097        let x = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
1098        let abs = x.abs().unwrap();
1099
1100        for p in [0.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
1101            let err = p_norm(&abs, &[0], p).unwrap_err();
1102            assert!(
1103                err.to_string().contains("finite") || err.to_string().contains("nonzero"),
1104                "expected finite nonzero order error, got {err:?}"
1105            );
1106        }
1107    }
1108}
tenferro_linalg/traced.rs

tenferro_linalg/
traced.rs
