tenferro_linalg/
eager_ext.rs

1use std::sync::Arc;
2
3use tenferro_ad::error::{Error, Result};
4use tenferro_ad::extension::apply_eager;
5use tenferro_ad::EagerTensor;
6
7use crate::extension::{
8    validate_derivative_eps, EighOptions, LinalgExtensionOp, LinalgOp, QrOptions, SvdOptions,
9};
10use crate::register_runtime;
11
12/// Linear algebra extension methods for [`EagerTensor`].
13pub trait EagerTensorLinalgExt {
14    fn svd(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)>;
15    fn svd_with_options(
16        &self,
17        options: SvdOptions,
18    ) -> Result<(EagerTensor, EagerTensor, EagerTensor)>;
19    fn qr(&self) -> Result<(EagerTensor, EagerTensor)>;
20    fn qr_with_options(&self, options: QrOptions) -> Result<(EagerTensor, EagerTensor)>;
21    fn lu(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor, EagerTensor)>;
22    fn full_piv_lu(
23        &self,
24    ) -> Result<(
25        EagerTensor,
26        EagerTensor,
27        EagerTensor,
28        EagerTensor,
29        EagerTensor,
30    )>;
31    fn full_piv_lu_solve(&self, b: &EagerTensor) -> Result<EagerTensor>;
32    fn solve(&self, b: &EagerTensor) -> Result<EagerTensor>;
33    fn cholesky(&self) -> Result<EagerTensor>;
34    fn eigh(&self) -> Result<(EagerTensor, EagerTensor)>;
35    fn eigh_with_options(&self, options: EighOptions) -> Result<(EagerTensor, EagerTensor)>;
36    fn eig(&self) -> Result<(EagerTensor, EagerTensor)>;
37    fn triangular_solve(
38        &self,
39        b: &EagerTensor,
40        left_side: bool,
41        lower: bool,
42        transpose_a: bool,
43        unit_diagonal: bool,
44    ) -> Result<EagerTensor>;
45}
46
47impl EagerTensorLinalgExt for EagerTensor {
48    fn svd(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
49        svd(self)
50    }
51
52    fn svd_with_options(
53        &self,
54        options: SvdOptions,
55    ) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
56        svd_with_options(self, options)
57    }
58
59    fn qr(&self) -> Result<(EagerTensor, EagerTensor)> {
60        qr(self)
61    }
62
63    fn qr_with_options(&self, options: QrOptions) -> Result<(EagerTensor, EagerTensor)> {
64        qr_with_options(self, options)
65    }
66
67    fn lu(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor, EagerTensor)> {
68        lu(self)
69    }
70
71    fn full_piv_lu(
72        &self,
73    ) -> Result<(
74        EagerTensor,
75        EagerTensor,
76        EagerTensor,
77        EagerTensor,
78        EagerTensor,
79    )> {
80        full_piv_lu(self)
81    }
82
83    fn full_piv_lu_solve(&self, b: &EagerTensor) -> Result<EagerTensor> {
84        full_piv_lu_solve(self, b)
85    }
86
87    fn solve(&self, b: &EagerTensor) -> Result<EagerTensor> {
88        solve(self, b)
89    }
90
91    fn cholesky(&self) -> Result<EagerTensor> {
92        cholesky(self)
93    }
94
95    fn eigh(&self) -> Result<(EagerTensor, EagerTensor)> {
96        eigh(self)
97    }
98
99    fn eigh_with_options(&self, options: EighOptions) -> Result<(EagerTensor, EagerTensor)> {
100        eigh_with_options(self, options)
101    }
102
103    fn eig(&self) -> Result<(EagerTensor, EagerTensor)> {
104        eig(self)
105    }
106
107    fn triangular_solve(
108        &self,
109        b: &EagerTensor,
110        left_side: bool,
111        lower: bool,
112        transpose_a: bool,
113        unit_diagonal: bool,
114    ) -> Result<EagerTensor> {
115        triangular_solve(self, b, left_side, lower, transpose_a, unit_diagonal)
116    }
117}
118
119fn apply_linalg_eager(op: LinalgOp, inputs: &[&EagerTensor]) -> Result<Vec<EagerTensor>> {
120    if let Some(first) = inputs.first() {
121        first
122            .runtime()
123            .register_extension(register_runtime)
124            .map_err(|err| Error::Internal(err.to_string()))?;
125    }
126    apply_eager(Arc::new(LinalgExtensionOp::new(op)), inputs)
127}
128
129/// Singular value decomposition for eager tensors.
130///
131/// # Examples
132///
133/// ```rust
134/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
135/// use tenferro_linalg::EagerTensorLinalgExt;
136///
137/// let ctx = EagerRuntime::new();
138/// let a = EagerTensor::from_tensor_in(
139///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0]).unwrap(),
140///     ctx,
141/// ).unwrap();
142/// let (_u, s, _vt) = a.svd()?;
143/// assert_eq!(s.shape(), &[2]);
144/// # Ok::<(), tenferro_ad::Error>(())
145/// ```
146pub fn svd(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
147    svd_with_options(a, SvdOptions::default())
148}
149
150/// Singular value decomposition for eager tensors with explicit options.
151///
152/// `derivative_eps` regularizes decomposition derivative formulas. It is not a
153/// backend SVD solver tolerance.
154///
155/// # Examples
156///
157/// ```rust
158/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
159/// use tenferro_linalg::{EagerTensorLinalgExt, SvdGauge, SvdOptions};
160///
161/// let ctx = EagerRuntime::new();
162/// let a = EagerTensor::from_tensor_in(
163///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0]).unwrap(),
164///     ctx,
165/// ).unwrap();
166/// let options = SvdOptions::default()
167///     .gauge(SvdGauge::CanonicalPivot)
168///     .derivative_eps(1.0e-10);
169/// let (_u, s, _vt) = a.svd_with_options(options)?;
170/// assert_eq!(s.shape(), &[2]);
171/// # Ok::<(), tenferro_ad::Error>(())
172/// ```
173pub fn svd_with_options(
174    a: &EagerTensor,
175    options: SvdOptions,
176) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
177    validate_derivative_eps("svd_with_options", options.derivative_eps)?;
178    let mut outputs = apply_linalg_eager(
179        LinalgOp::Svd {
180            derivative_eps: options.derivative_eps,
181            gauge: options.gauge,
182        },
183        &[a],
184    )?
185    .into_iter();
186    match (
187        outputs.next(),
188        outputs.next(),
189        outputs.next(),
190        outputs.next(),
191    ) {
192        (Some(u), Some(s), Some(vt), None) => Ok((u, s, vt)),
193        _ => Err(Error::Internal(
194            "svd eager op returned an unexpected number of outputs".to_string(),
195        )),
196    }
197}
198
199/// QR decomposition for eager tensors.
200///
201/// # Examples
202///
203/// ```rust
204/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
205/// use tenferro_linalg::EagerTensorLinalgExt;
206///
207/// let ctx = EagerRuntime::new();
208/// let a = EagerTensor::from_tensor_in(
209///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap(),
210///     ctx,
211/// ).unwrap();
212/// let (q, r) = a.qr()?;
213/// assert_eq!(q.shape(), &[2, 2]);
214/// assert_eq!(r.shape(), &[2, 2]);
215/// # Ok::<(), tenferro_ad::Error>(())
216/// ```
217pub fn qr(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor)> {
218    qr_with_options(a, QrOptions::default())
219}
220
221/// QR decomposition for eager tensors with explicit options.
222///
223/// `gauge` controls optional sign or phase post-processing.
224///
225/// # Examples
226///
227/// ```rust
228/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
229/// use tenferro_linalg::{EagerTensorLinalgExt, QrGauge, QrOptions};
230///
231/// let ctx = EagerRuntime::new();
232/// let a = EagerTensor::from_tensor_in(
233///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap(),
234///     ctx,
235/// ).unwrap();
236/// let (q, r) = a.qr_with_options(QrOptions::default().gauge(QrGauge::PositiveDiagonal))?;
237/// assert_eq!(q.shape(), &[2, 2]);
238/// assert_eq!(r.shape(), &[2, 2]);
239/// # Ok::<(), tenferro_ad::Error>(())
240/// ```
241pub fn qr_with_options(a: &EagerTensor, options: QrOptions) -> Result<(EagerTensor, EagerTensor)> {
242    two_outputs(
243        apply_linalg_eager(
244            LinalgOp::Qr {
245                gauge: options.gauge,
246            },
247            &[a],
248        )?,
249        "qr",
250    )
251}
252
253/// LU factorization for eager tensors.
254///
255/// # Examples
256///
257/// ```rust
258/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
259/// use tenferro_linalg::EagerTensorLinalgExt;
260///
261/// let ctx = EagerRuntime::new();
262/// let a = EagerTensor::from_tensor_in(
263///     Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f64, 1.0, 1.0, 0.0]).unwrap(),
264///     ctx,
265/// ).unwrap();
266/// let (_p, l, u, parity) = a.lu()?;
267/// assert_eq!(l.shape(), &[2, 2]);
268/// assert_eq!(u.shape(), &[2, 2]);
269/// assert_eq!(parity.shape(), &[] as &[usize]);
270/// # Ok::<(), tenferro_ad::Error>(())
271/// ```
272pub fn lu(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor, EagerTensor, EagerTensor)> {
273    let mut outputs = apply_linalg_eager(LinalgOp::Lu, &[a])?.into_iter();
274    match (
275        outputs.next(),
276        outputs.next(),
277        outputs.next(),
278        outputs.next(),
279        outputs.next(),
280    ) {
281        (Some(p), Some(l), Some(u), Some(parity), None) => Ok((p, l, u, parity)),
282        _ => Err(Error::Internal(
283            "lu eager op returned an unexpected number of outputs".to_string(),
284        )),
285    }
286}
287
288/// Complete-pivot LU factorization for eager tensors.
289///
290/// Returns `(P, L, U, Q, parity)` with reconstruction convention
291/// `A = P^T * L * U * Q`, equivalently `P * A * Q^T = L * U`. `parity` is a
292/// scalar real tensor containing `+1` or `-1`: `F32` for `F32`/`C32` inputs and
293/// `F64` for `F64`/`C64` inputs.
294///
295/// # Examples
296///
297/// ```rust
298/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
299/// use tenferro_linalg::EagerTensorLinalgExt;
300///
301/// let ctx = EagerRuntime::new();
302/// let a = EagerTensor::from_tensor_in(
303///     Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f64, 2.0, 1.0, 3.0]).unwrap(),
304///     ctx,
305/// ).unwrap();
306/// let (p, _l, _u, q, parity) = a.full_piv_lu()?;
307/// assert_eq!(p.shape(), &[2, 2]);
308/// assert_eq!(q.shape(), &[2, 2]);
309/// assert_eq!(parity.shape(), &[] as &[usize]);
310/// # Ok::<(), tenferro_ad::Error>(())
311/// ```
312pub fn full_piv_lu(
313    a: &EagerTensor,
314) -> Result<(
315    EagerTensor,
316    EagerTensor,
317    EagerTensor,
318    EagerTensor,
319    EagerTensor,
320)> {
321    let mut outputs = apply_linalg_eager(LinalgOp::FullPivLu, &[a])?.into_iter();
322    match (
323        outputs.next(),
324        outputs.next(),
325        outputs.next(),
326        outputs.next(),
327        outputs.next(),
328        outputs.next(),
329    ) {
330        (Some(p), Some(l), Some(u), Some(q), Some(parity), None) => Ok((p, l, u, q, parity)),
331        _ => Err(Error::Internal(
332            "full_piv_lu eager op returned an unexpected number of outputs".to_string(),
333        )),
334    }
335}
336
337/// Solve a linear system using complete-pivot LU behavior.
338///
339/// # Examples
340///
341/// ```rust
342/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
343/// use tenferro_linalg::EagerTensorLinalgExt;
344///
345/// let ctx = EagerRuntime::new();
346/// let a = EagerTensor::from_tensor_in(
347///     Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f64, 2.0, 1.0, 3.0]).unwrap(),
348///     ctx.clone(),
349/// ).unwrap();
350/// let b = EagerTensor::from_tensor_in(
351///     Tensor::from_vec_col_major(vec![2, 1], vec![-1.0_f64, 5.0]).unwrap(),
352///     ctx,
353/// ).unwrap();
354/// let x = a.full_piv_lu_solve(&b)?;
355/// assert_eq!(x.shape(), &[2, 1]);
356/// # Ok::<(), tenferro_ad::Error>(())
357/// ```
358pub fn full_piv_lu_solve(a: &EagerTensor, b: &EagerTensor) -> Result<EagerTensor> {
359    one_output(
360        apply_linalg_eager(LinalgOp::FullPivLuSolve { transpose_a: false }, &[a, b])?,
361        "full_piv_lu_solve",
362    )
363}
364
365/// Solve a linear system for eager tensors.
366///
367/// # Examples
368///
369/// ```rust
370/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
371/// use tenferro_linalg::EagerTensorLinalgExt;
372///
373/// let ctx = EagerRuntime::new();
374/// let a = EagerTensor::from_tensor_in(
375///     Tensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 4.0]).unwrap(),
376///     ctx.clone(),
377/// ).unwrap();
378/// let b = EagerTensor::from_tensor_in(
379///     Tensor::from_vec_col_major(vec![2, 1], vec![4.0_f64, 8.0]).unwrap(),
380///     ctx,
381/// ).unwrap();
382/// let x = a.solve(&b)?;
383/// assert_eq!(x.shape(), &[2, 1]);
384/// # Ok::<(), tenferro_ad::Error>(())
385/// ```
386pub fn solve(a: &EagerTensor, b: &EagerTensor) -> Result<EagerTensor> {
387    let mut factor_outputs = apply_linalg_eager(LinalgOp::LuFactor, &[a])?.into_iter();
388    let (packed_lu, pivots) = match (
389        factor_outputs.next(),
390        factor_outputs.next(),
391        factor_outputs.next(),
392        factor_outputs.next(),
393    ) {
394        (Some(packed_lu), Some(pivots), Some(_parity), None) => (packed_lu, pivots),
395        _ => {
396            return Err(Error::Internal(
397                "lu_factor eager op returned an unexpected number of outputs".to_string(),
398            ));
399        }
400    };
401    one_output(
402        apply_linalg_eager(
403            LinalgOp::LuSolvePrepared {
404                transpose_a: false,
405                conjugate_a: false,
406            },
407            &[a, &packed_lu, &pivots, b],
408        )?,
409        "solve",
410    )
411}
412
413/// Cholesky factorization for eager tensors.
414///
415/// # Examples
416///
417/// ```rust
418/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
419/// use tenferro_linalg::EagerTensorLinalgExt;
420///
421/// let ctx = EagerRuntime::new();
422/// let a = EagerTensor::from_tensor_in(
423///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap(),
424///     ctx,
425/// ).unwrap();
426/// let l = a.cholesky()?;
427/// assert_eq!(l.shape(), &[2, 2]);
428/// # Ok::<(), tenferro_ad::Error>(())
429/// ```
430pub fn cholesky(a: &EagerTensor) -> Result<EagerTensor> {
431    one_output(apply_linalg_eager(LinalgOp::Cholesky, &[a])?, "cholesky")
432}
433
434/// Hermitian eigenvalue decomposition for eager tensors.
435///
436/// # Examples
437///
438/// ```rust
439/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
440/// use tenferro_linalg::EagerTensorLinalgExt;
441///
442/// let ctx = EagerRuntime::new();
443/// let a = EagerTensor::from_tensor_in(
444///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 3.0]).unwrap(),
445///     ctx,
446/// ).unwrap();
447/// let (values, vectors) = a.eigh()?;
448/// assert_eq!(values.shape(), &[2]);
449/// assert_eq!(vectors.shape(), &[2, 2]);
450/// # Ok::<(), tenferro_ad::Error>(())
451/// ```
452pub fn eigh(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor)> {
453    eigh_with_options(a, EighOptions::default())
454}
455
456/// Hermitian eigenvalue decomposition for eager tensors with explicit options.
457///
458/// `derivative_eps` regularizes derivative formulas for repeated or nearly
459/// repeated eigenvalues. It is not a backend eigensolver tolerance.
460///
461/// # Examples
462///
463/// ```rust
464/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
465/// use tenferro_linalg::{EagerTensorLinalgExt, EighGauge, EighOptions};
466///
467/// let ctx = EagerRuntime::new();
468/// let a = EagerTensor::from_tensor_in(
469///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 3.0]).unwrap(),
470///     ctx,
471/// ).unwrap();
472/// let (values, vectors) = a
473///     .eigh_with_options(
474///         EighOptions::default()
475///             .gauge(EighGauge::CanonicalPivot)
476///             .derivative_eps(1.0e-10),
477///     )?;
478/// assert_eq!(values.shape(), &[2]);
479/// assert_eq!(vectors.shape(), &[2, 2]);
480/// # Ok::<(), tenferro_ad::Error>(())
481/// ```
482pub fn eigh_with_options(
483    a: &EagerTensor,
484    options: EighOptions,
485) -> Result<(EagerTensor, EagerTensor)> {
486    validate_derivative_eps("eigh_with_options", options.derivative_eps)?;
487    two_outputs(
488        apply_linalg_eager(
489            LinalgOp::Eigh {
490                derivative_eps: options.derivative_eps,
491                gauge: options.gauge,
492            },
493            &[a],
494        )?,
495        "eigh",
496    )
497}
498
499/// General eigenvalue decomposition for eager tensors.
500///
501/// # Examples
502///
503/// ```rust
504/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
505/// use tenferro_linalg::EagerTensorLinalgExt;
506///
507/// let ctx = EagerRuntime::new();
508/// let a = EagerTensor::from_tensor_in(
509///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 3.0]).unwrap(),
510///     ctx,
511/// ).unwrap();
512/// let (values, vectors) = a.eig()?;
513/// assert_eq!(values.shape(), &[2]);
514/// assert_eq!(vectors.shape(), &[2, 2]);
515/// # Ok::<(), tenferro_ad::Error>(())
516/// ```
517pub fn eig(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor)> {
518    two_outputs(
519        apply_linalg_eager(
520            LinalgOp::Eig {
521                input_dtype: a.dtype(),
522            },
523            &[a],
524        )?,
525        "eig",
526    )
527}
528
529/// Triangular solve for eager tensors.
530///
531/// # Examples
532///
533/// ```rust
534/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
535/// use tenferro_linalg::EagerTensorLinalgExt;
536///
537/// let ctx = EagerRuntime::new();
538/// let a = EagerTensor::from_tensor_in(
539///     Tensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 1.0, 3.0]).unwrap(),
540///     ctx.clone(),
541/// ).unwrap();
542/// let b = EagerTensor::from_tensor_in(
543///     Tensor::from_vec_col_major(vec![2, 1], vec![2.0_f64, 7.0]).unwrap(),
544///     ctx,
545/// ).unwrap();
546/// let x = a.triangular_solve(&b, true, true, false, false)?;
547/// assert_eq!(x.shape(), &[2, 1]);
548/// # Ok::<(), tenferro_ad::Error>(())
549/// ```
550pub fn triangular_solve(
551    a: &EagerTensor,
552    b: &EagerTensor,
553    left_side: bool,
554    lower: bool,
555    transpose_a: bool,
556    unit_diagonal: bool,
557) -> Result<EagerTensor> {
558    one_output(
559        apply_linalg_eager(
560            LinalgOp::TriangularSolve {
561                left_side,
562                lower,
563                transpose_a,
564                unit_diagonal,
565            },
566            &[a, b],
567        )?,
568        "triangular_solve",
569    )
570}
571
572fn one_output(outputs: Vec<EagerTensor>, name: &str) -> Result<EagerTensor> {
573    let mut outputs = outputs.into_iter();
574    match (outputs.next(), outputs.next()) {
575        (Some(output), None) => Ok(output),
576        _ => Err(Error::Internal(format!(
577            "{name} eager op returned an unexpected number of outputs"
578        ))),
579    }
580}
581
582fn two_outputs(outputs: Vec<EagerTensor>, name: &str) -> Result<(EagerTensor, EagerTensor)> {
583    let mut outputs = outputs.into_iter();
584    match (outputs.next(), outputs.next(), outputs.next()) {
585        (Some(lhs), Some(rhs), None) => Ok((lhs, rhs)),
586        _ => Err(Error::Internal(format!(
587            "{name} eager op returned an unexpected number of outputs"
588        ))),
589    }
590}
tenferro_linalg/eager_ext.rs

tenferro_linalg/
eager_ext.rs
