tenferro_linalg/
eager_ext.rs

1use std::sync::{Arc, OnceLock};
2
3use tenferro_ad::error::{Error, Result};
4use tenferro_ad::extension::apply_eager_with_extension_session;
5use tenferro_ad::EagerTensor;
6use tenferro_runtime::{ErrorPhase, ExtensionModule};
7
8use crate::eager_composites;
9use crate::extension::{
10    execute_linalg_extension_reads, extension_module, validate_derivative_eps, EighOptions,
11    LinalgExtensionOp, LinalgOp, QrOptions, SvdOptions,
12};
13
14/// Linear algebra extension methods for [`EagerTensor`].
15pub trait EagerTensorLinalgExt {
16    /// # Errors
17    ///
18    /// Returns `Error::Validation` for an invalid rank, shape, or dtype,
19    /// `Error::Extension` with an unsupported-operation or unsupported-dtype
20    /// source when the selected backend cannot execute the decomposition, and
21    /// `Error::RuntimeState` when the eager runtime or backend is unavailable.
22    fn svd(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)>;
23    /// # Errors
24    ///
25    /// Returns `Error::Validation` when `derivative_eps` is non-finite or
26    /// non-positive, `Error::Extension` for unsupported dtypes or numerical
27    /// non-convergence, and `Error::Internal` if the extension violates its
28    /// output-count contract.
29    fn svd_with_options(
30        &self,
31        options: SvdOptions,
32    ) -> Result<(EagerTensor, EagerTensor, EagerTensor)>;
33    /// Full-matrices SVD returning square `U (m x m)` and `Vh (n x n)`, whose
34    /// trailing `n - rank` rows span the input's right nullspace.
35    ///
36    /// # Errors
37    ///
38    /// Returns `Error::Validation` for an invalid rank, and `Error::Extension`
39    /// with an unsupported-operation source when the active backend does not
40    /// implement full-matrices SVD (the CPU faer provider supports it; the
41    /// LAPACK provider and GPU backends return an unsupported error in this
42    /// slice). Automatic differentiation through the full variant is
43    /// unsupported and surfaces a typed error rather than a silent thin
44    /// fallback.
45    fn svd_full(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)>;
46    /// # Errors
47    ///
48    /// Returns `Error::Validation` for invalid matrix rank or shape,
49    /// `Error::Extension` for unsupported dtypes or numerical failure, and
50    /// `Error::RuntimeState` when the eager runtime or backend is unavailable.
51    fn qr(&self) -> Result<(EagerTensor, EagerTensor)>;
52    /// # Errors
53    ///
54    /// Returns `Error::Validation` for invalid matrix rank or shape,
55    /// `Error::Extension` for unsupported dtypes or numerical failure, and
56    /// `Error::Internal` if the extension violates its output-count contract.
57    fn qr_with_options(&self, options: QrOptions) -> Result<(EagerTensor, EagerTensor)>;
58    /// # Errors
59    ///
60    /// Returns `Error::Validation` for an invalid matrix rank or shape,
61    /// `Error::Extension` for an unsupported dtype or singular numerical
62    /// result, and `Error::RuntimeState` when execution cannot access its
63    /// backend.
64    fn lu(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor, EagerTensor)>;
65    /// # Errors
66    ///
67    /// Returns `Error::Validation` for an invalid matrix rank or shape,
68    /// `Error::Extension` for unsupported dtypes or singular numerical
69    /// results, and `Error::Internal` if the extension violates its output
70    /// contract.
71    fn full_piv_lu(
72        &self,
73    ) -> Result<(
74        EagerTensor,
75        EagerTensor,
76        EagerTensor,
77        EagerTensor,
78        EagerTensor,
79    )>;
80    /// # Errors
81    ///
82    /// Returns `Error::Validation` when `a` and `b` have incompatible matrix
83    /// or batch shapes, `Error::Extension` for an unsupported dtype or
84    /// singular system, and `Error::RuntimeState` when the backend is
85    /// unavailable.
86    fn full_piv_lu_solve(&self, b: &EagerTensor) -> Result<EagerTensor>;
87    /// # Errors
88    ///
89    /// Returns `Error::Validation` for incompatible matrix, batch, or dtype
90    /// metadata, `Error::Extension` for an unsupported dtype or singular
91    /// system, and `Error::RuntimeState` when the backend is unavailable.
92    fn solve(&self, b: &EagerTensor) -> Result<EagerTensor>;
93    /// Least-squares solve `argmin_x ||A x - b||_2` for a tall or square,
94    /// full-column-rank `A`, via the thin QR factorization.
95    ///
96    /// # Errors
97    ///
98    /// Returns `Error::Validation` when `A` or `b` is not a matrix (rank
99    /// `>= 2`), when `A` is wide (`rows < cols`; underdetermined), or when the
100    /// dtype is not floating-point or complex; `Error::Extension` for backend
101    /// QR or triangular-solve failures; and `Error::RuntimeState` when the
102    /// backend is unavailable. Rank-deficient `A` is not detected and yields an
103    /// ill-defined result.
104    fn lstsq(&self, b: &EagerTensor) -> Result<EagerTensor>;
105    /// # Errors
106    ///
107    /// Returns `Error::Validation` for a non-square or invalid-rank input,
108    /// `Error::Extension` for unsupported dtypes or a non-positive-definite
109    /// matrix, and `Error::RuntimeState` when the backend is unavailable.
110    fn cholesky(&self) -> Result<EagerTensor>;
111    /// # Errors
112    ///
113    /// Returns `Error::Validation` for a non-square or invalid-rank input,
114    /// `Error::Extension` for unsupported dtypes or numerical non-convergence,
115    /// and `Error::RuntimeState` when the backend is unavailable.
116    fn eigh(&self) -> Result<(EagerTensor, EagerTensor)>;
117    /// # Errors
118    ///
119    /// Returns `Error::Validation` for an invalid rank, shape, or
120    /// `derivative_eps`, `Error::Extension` for unsupported dtypes or
121    /// non-convergence, and `Error::Internal` for an output-count violation.
122    fn eigh_with_options(&self, options: EighOptions) -> Result<(EagerTensor, EagerTensor)>;
123    /// # Errors
124    ///
125    /// Returns `Error::Validation` for a non-square or invalid-rank input,
126    /// `Error::Extension` for unsupported dtypes or numerical non-convergence,
127    /// and `Error::RuntimeState` when the backend is unavailable.
128    fn eig(&self) -> Result<(EagerTensor, EagerTensor)>;
129    /// # Errors
130    ///
131    /// Returns `Error::Validation` for incompatible matrix, batch, or dtype
132    /// metadata, `Error::Extension` for unsupported dtypes or a singular
133    /// system, and `Error::RuntimeState` when the backend is unavailable.
134    fn triangular_solve(
135        &self,
136        b: &EagerTensor,
137        left_side: bool,
138        lower: bool,
139        transpose_a: bool,
140        unit_diagonal: bool,
141    ) -> Result<EagerTensor>;
142    /// Return the determinant sign and logarithm of its absolute value.
143    ///
144    /// # Errors
145    ///
146    /// Returns `Error::Validation` for a non-square input, `Error::Extension`
147    /// for an unsupported dtype or numerical failure, and `Error::Internal`
148    /// if a primitive violates its output contract.
149    fn slogdet(&self) -> Result<(EagerTensor, EagerTensor)>;
150    /// Return the determinant.
151    ///
152    /// # Errors
153    ///
154    /// Returns `Error::Validation`, `Error::Extension`, `Error::RuntimeState`,
155    /// or `Error::Internal` under the conditions documented by [`Self::slogdet`].
156    fn det(&self) -> Result<EagerTensor>;
157    /// Return the matrix inverse.
158    ///
159    /// # Errors
160    ///
161    /// Returns `Error::Validation` for invalid matrix metadata,
162    /// `Error::Extension` for an unsupported dtype or singular solve, and
163    /// `Error::RuntimeState` when eager execution cannot access its backend.
164    fn inv(&self) -> Result<EagerTensor>;
165    /// Return Hermitian eigenvalues without eigenvectors.
166    ///
167    /// # Errors
168    ///
169    /// Returns the validation, unsupported-dtype, numerical-convergence,
170    /// runtime-state, or output-contract errors reported by [`Self::eigh`].
171    fn eigvalsh(&self) -> Result<EagerTensor>;
172    /// Return general eigenvalues without eigenvectors.
173    ///
174    /// # Errors
175    ///
176    /// Returns the validation, unsupported-dtype, numerical-convergence,
177    /// runtime-state, or output-contract errors reported by [`Self::eig`].
178    fn eigvals(&self) -> Result<EagerTensor>;
179    /// Return the Moore-Penrose pseudoinverse with the default tolerance.
180    ///
181    /// # Errors
182    ///
183    /// Returns `Error::Validation` for invalid matrix metadata,
184    /// `Error::Extension` for unsupported SVD execution, and
185    /// `Error::Internal` for an unexpected decomposition output contract.
186    fn pinv(&self) -> Result<EagerTensor>;
187    /// Return the Moore-Penrose pseudoinverse with an explicit relative tolerance.
188    ///
189    /// # Errors
190    ///
191    /// Returns `Error::Validation` when `rtol` is negative or non-finite, plus
192    /// the `Error::Extension` and output-contract errors reported by [`Self::pinv`].
193    fn pinv_with_rtol(&self, rtol: f64) -> Result<EagerTensor>;
194    /// Return a vector, matrix, or tensor norm.
195    ///
196    /// # Errors
197    ///
198    /// Returns `Error::Validation` for an unsupported order/dimension
199    /// combination or invalid axes, `Error::Extension` when a required matrix
200    /// decomposition is unsupported, and eager runtime/backend failures.
201    fn norm(&self, ord: Option<f64>, dim: Option<&[usize]>, keepdim: bool) -> Result<EagerTensor>;
202}
203
204impl EagerTensorLinalgExt for EagerTensor {
205    fn svd(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
206        svd(self)
207    }
208
209    fn svd_with_options(
210        &self,
211        options: SvdOptions,
212    ) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
213        svd_with_options(self, options)
214    }
215
216    fn svd_full(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
217        svd_full(self)
218    }
219
220    fn qr(&self) -> Result<(EagerTensor, EagerTensor)> {
221        qr(self)
222    }
223
224    fn qr_with_options(&self, options: QrOptions) -> Result<(EagerTensor, EagerTensor)> {
225        qr_with_options(self, options)
226    }
227
228    fn lu(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor, EagerTensor)> {
229        lu(self)
230    }
231
232    fn full_piv_lu(
233        &self,
234    ) -> Result<(
235        EagerTensor,
236        EagerTensor,
237        EagerTensor,
238        EagerTensor,
239        EagerTensor,
240    )> {
241        full_piv_lu(self)
242    }
243
244    fn full_piv_lu_solve(&self, b: &EagerTensor) -> Result<EagerTensor> {
245        full_piv_lu_solve(self, b)
246    }
247
248    fn solve(&self, b: &EagerTensor) -> Result<EagerTensor> {
249        solve(self, b)
250    }
251
252    fn lstsq(&self, b: &EagerTensor) -> Result<EagerTensor> {
253        eager_composites::lstsq(self, b)
254    }
255
256    fn cholesky(&self) -> Result<EagerTensor> {
257        cholesky(self)
258    }
259
260    fn eigh(&self) -> Result<(EagerTensor, EagerTensor)> {
261        eigh(self)
262    }
263
264    fn eigh_with_options(&self, options: EighOptions) -> Result<(EagerTensor, EagerTensor)> {
265        eigh_with_options(self, options)
266    }
267
268    fn eig(&self) -> Result<(EagerTensor, EagerTensor)> {
269        eig(self)
270    }
271
272    fn triangular_solve(
273        &self,
274        b: &EagerTensor,
275        left_side: bool,
276        lower: bool,
277        transpose_a: bool,
278        unit_diagonal: bool,
279    ) -> Result<EagerTensor> {
280        triangular_solve(self, b, left_side, lower, transpose_a, unit_diagonal)
281    }
282
283    fn slogdet(&self) -> Result<(EagerTensor, EagerTensor)> {
284        eager_composites::slogdet(self)
285    }
286
287    fn det(&self) -> Result<EagerTensor> {
288        eager_composites::det(self)
289    }
290
291    fn inv(&self) -> Result<EagerTensor> {
292        eager_composites::inv(self)
293    }
294
295    fn eigvalsh(&self) -> Result<EagerTensor> {
296        eager_composites::eigvalsh(self)
297    }
298
299    fn eigvals(&self) -> Result<EagerTensor> {
300        eager_composites::eigvals(self)
301    }
302
303    fn pinv(&self) -> Result<EagerTensor> {
304        eager_composites::pinv(self)
305    }
306
307    fn pinv_with_rtol(&self, rtol: f64) -> Result<EagerTensor> {
308        eager_composites::pinv_with_rtol(self, rtol)
309    }
310
311    fn norm(&self, ord: Option<f64>, dim: Option<&[usize]>, keepdim: bool) -> Result<EagerTensor> {
312        eager_composites::norm(self, ord, dim, keepdim)
313    }
314}
315
316fn apply_linalg_eager(op: LinalgOp, inputs: &[&EagerTensor]) -> Result<Vec<EagerTensor>> {
317    let op = Arc::new(LinalgExtensionOp::new(op));
318    let execute_op = Arc::clone(&op);
319    let module = eager_cpu_extension_module()?;
320    apply_eager_with_extension_session(op, inputs, module, move |_op, input_reads, ctx| {
321        execute_linalg_extension_reads(&execute_op, input_reads, ctx)
322    })
323}
324
325fn eager_cpu_extension_module() -> Result<Arc<dyn ExtensionModule>> {
326    static MODULE: OnceLock<Arc<dyn ExtensionModule>> = OnceLock::new();
327    if let Some(module) = MODULE.get() {
328        return Ok(Arc::clone(module));
329    }
330
331    let engine_id = tenferro_cpu::runtime_engine_id().map_err(eager_runtime_config_error)?;
332    let module = extension_module::<tenferro_cpu::CpuBackend>(engine_id)
333        .map_err(eager_runtime_config_error)?;
334    let _ = MODULE.set(Arc::clone(&module));
335    Ok(MODULE.get().cloned().unwrap_or(module))
336}
337
338fn eager_runtime_config_error(source: tenferro_runtime::RuntimeConfigError) -> Error {
339    Error::runtime_state_source(
340        "tenferro_linalg::eager_extension_module",
341        ErrorPhase::Execution,
342        source,
343    )
344}
345
346/// Singular value decomposition for eager tensors.
347///
348/// # Examples
349///
350/// ```rust
351/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
352/// use tenferro_linalg::EagerTensorLinalgExt;
353///
354/// let ctx = EagerRuntime::new()?;
355/// let a = EagerTensor::from_tensor_in(
356///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0]).unwrap(),
357///     ctx,
358/// ).unwrap();
359/// let (_u, s, _vt) = a.svd()?;
360/// assert_eq!(s.shape(), &[2]);
361/// # Ok::<(), tenferro_ad::Error>(())
362/// ```
363///
364/// # Errors
365///
366/// Returns `Error::Validation` for an invalid rank, matrix shape, or dtype,
367/// `Error::Extension` with an unsupported-dtype or non-convergence source when
368/// the backend cannot compute the decomposition, and `Error::RuntimeState`
369/// when the eager runtime or backend is unavailable.
370pub fn svd(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
371    svd_with_options(a, SvdOptions::default())
372}
373
374/// Singular value decomposition for eager tensors with explicit options.
375///
376/// `derivative_eps` regularizes decomposition derivative formulas. It is not a
377/// backend SVD solver tolerance.
378///
379/// # Examples
380///
381/// ```rust
382/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
383/// use tenferro_linalg::{EagerTensorLinalgExt, SvdGauge, SvdOptions};
384///
385/// let ctx = EagerRuntime::new()?;
386/// let a = EagerTensor::from_tensor_in(
387///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0]).unwrap(),
388///     ctx,
389/// ).unwrap();
390/// let options = SvdOptions::default()
391///     .gauge(SvdGauge::CanonicalPivot)
392///     .derivative_eps(1.0e-10);
393/// let (_u, s, _vt) = a.svd_with_options(options)?;
394/// assert_eq!(s.shape(), &[2]);
395/// # Ok::<(), tenferro_ad::Error>(())
396/// ```
397///
398/// # Errors
399///
400/// Returns `Error::Validation` when `derivative_eps` is non-finite or
401/// non-positive, `Error::Extension` for unsupported dtypes or numerical
402/// non-convergence, and `Error::Internal` if the extension returns an
403/// unexpected number of outputs.
404pub fn svd_with_options(
405    a: &EagerTensor,
406    options: SvdOptions,
407) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
408    validate_derivative_eps("svd_with_options", options.derivative_eps)?;
409    let mut outputs = apply_linalg_eager(
410        LinalgOp::Svd {
411            derivative_eps: options.derivative_eps,
412            gauge: options.gauge,
413        },
414        &[a],
415    )?
416    .into_iter();
417    match (
418        outputs.next(),
419        outputs.next(),
420        outputs.next(),
421        outputs.next(),
422    ) {
423        (Some(u), Some(s), Some(vt), None) => Ok((u, s, vt)),
424        _ => Err(Error::Internal(
425            "svd eager op returned an unexpected number of outputs".to_string(),
426        )),
427    }
428}
429
430/// Full-matrices singular value decomposition for eager tensors.
431///
432/// Returns `(U, S, Vh)` with square `U` (`m x m`) and `Vh` (`n x n`); `S` holds
433/// `min(m, n)` singular values. The trailing `n - rank` rows of `Vh` span the
434/// input's right nullspace, which the thin [`svd`] cannot represent.
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![1, 2], vec![1.0_f64, 1.0]).unwrap(),
445///     ctx,
446/// ).unwrap();
447/// // Full SVD is provided by the CPU faer backend; other providers (e.g. the
448/// // LAPACK provider) report it as unsupported, so guard on the result.
449/// if let Ok((u, s, vh)) = a.svd_full() {
450///     assert_eq!(u.shape(), &[1, 1]);
451///     assert_eq!(s.shape(), &[1]);
452///     assert_eq!(vh.shape(), &[2, 2]);
453/// }
454/// # Ok::<(), tenferro_ad::Error>(())
455/// ```
456///
457/// # Errors
458///
459/// Returns `Error::Validation` for an invalid rank and `Error::Extension` with
460/// an unsupported-operation source when the active backend does not implement
461/// full-matrices SVD (the CPU faer provider supports it; the LAPACK provider
462/// and GPU backends are unsupported in this slice). AD through the full variant
463/// is unsupported and surfaces a typed error, not a silent thin fallback.
464pub fn svd_full(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
465    let mut outputs = apply_linalg_eager(LinalgOp::SvdFull, &[a])?.into_iter();
466    match (
467        outputs.next(),
468        outputs.next(),
469        outputs.next(),
470        outputs.next(),
471    ) {
472        (Some(u), Some(s), Some(vh), None) => Ok((u, s, vh)),
473        _ => Err(Error::Internal(
474            "svd_full eager op returned an unexpected number of outputs".to_string(),
475        )),
476    }
477}
478
479/// QR decomposition for eager tensors.
480///
481/// # Examples
482///
483/// ```rust
484/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
485/// use tenferro_linalg::EagerTensorLinalgExt;
486///
487/// let ctx = EagerRuntime::new()?;
488/// let a = EagerTensor::from_tensor_in(
489///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap(),
490///     ctx,
491/// ).unwrap();
492/// let (q, r) = a.qr()?;
493/// assert_eq!(q.shape(), &[2, 2]);
494/// assert_eq!(r.shape(), &[2, 2]);
495/// # Ok::<(), tenferro_ad::Error>(())
496/// ```
497///
498/// # Errors
499///
500/// Returns `Error::Validation` for an invalid rank or matrix shape,
501/// `Error::Extension` for an unsupported dtype or numerical failure, and
502/// `Error::RuntimeState` when the eager runtime or backend is unavailable.
503pub fn qr(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor)> {
504    qr_with_options(a, QrOptions::default())
505}
506
507/// QR decomposition for eager tensors with explicit options.
508///
509/// `gauge` controls optional sign or phase post-processing.
510///
511/// # Examples
512///
513/// ```rust
514/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
515/// use tenferro_linalg::{EagerTensorLinalgExt, QrGauge, QrOptions};
516///
517/// let ctx = EagerRuntime::new()?;
518/// let a = EagerTensor::from_tensor_in(
519///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap(),
520///     ctx,
521/// ).unwrap();
522/// let (q, r) = a.qr_with_options(QrOptions::default().gauge(QrGauge::PositiveDiagonal))?;
523/// assert_eq!(q.shape(), &[2, 2]);
524/// assert_eq!(r.shape(), &[2, 2]);
525/// # Ok::<(), tenferro_ad::Error>(())
526/// ```
527///
528/// # Errors
529///
530/// Returns `Error::Validation` for an invalid rank or matrix shape,
531/// `Error::Extension` for an unsupported dtype or numerical failure, and
532/// `Error::Internal` if the extension returns an unexpected number of outputs.
533pub fn qr_with_options(a: &EagerTensor, options: QrOptions) -> Result<(EagerTensor, EagerTensor)> {
534    two_outputs(
535        apply_linalg_eager(
536            LinalgOp::Qr {
537                gauge: options.gauge,
538            },
539            &[a],
540        )?,
541        "qr",
542    )
543}
544
545/// LU factorization for eager tensors.
546///
547/// # Examples
548///
549/// ```rust
550/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
551/// use tenferro_linalg::EagerTensorLinalgExt;
552///
553/// let ctx = EagerRuntime::new()?;
554/// let a = EagerTensor::from_tensor_in(
555///     Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f64, 1.0, 1.0, 0.0]).unwrap(),
556///     ctx,
557/// ).unwrap();
558/// let (_p, l, u, parity) = a.lu()?;
559/// assert_eq!(l.shape(), &[2, 2]);
560/// assert_eq!(u.shape(), &[2, 2]);
561/// assert_eq!(parity.shape(), &[] as &[usize]);
562/// # Ok::<(), tenferro_ad::Error>(())
563/// ```
564///
565/// # Errors
566///
567/// Returns `Error::Validation` for an invalid rank or matrix shape,
568/// `Error::Extension` for an unsupported dtype or singular numerical result,
569/// and `Error::RuntimeState` when the eager runtime or backend is unavailable.
570pub fn lu(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor, EagerTensor, EagerTensor)> {
571    let mut outputs = apply_linalg_eager(LinalgOp::Lu, &[a])?.into_iter();
572    match (
573        outputs.next(),
574        outputs.next(),
575        outputs.next(),
576        outputs.next(),
577        outputs.next(),
578    ) {
579        (Some(p), Some(l), Some(u), Some(parity), None) => Ok((p, l, u, parity)),
580        _ => Err(Error::Internal(
581            "lu eager op returned an unexpected number of outputs".to_string(),
582        )),
583    }
584}
585
586/// Complete-pivot LU factorization for eager tensors.
587///
588/// Returns `(P, L, U, Q, parity)` with reconstruction convention
589/// `A = P^T * L * U * Q`, equivalently `P * A * Q^T = L * U`. `parity` is a
590/// scalar real tensor containing `+1` or `-1`: `F32` for `F32`/`C32` inputs and
591/// `F64` for `F64`/`C64` inputs.
592///
593/// # Examples
594///
595/// ```rust
596/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
597/// use tenferro_linalg::EagerTensorLinalgExt;
598///
599/// let ctx = EagerRuntime::new()?;
600/// let a = EagerTensor::from_tensor_in(
601///     Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f64, 2.0, 1.0, 3.0]).unwrap(),
602///     ctx,
603/// ).unwrap();
604/// let (p, _l, _u, q, parity) = a.full_piv_lu()?;
605/// assert_eq!(p.shape(), &[2, 2]);
606/// assert_eq!(q.shape(), &[2, 2]);
607/// assert_eq!(parity.shape(), &[] as &[usize]);
608/// # Ok::<(), tenferro_ad::Error>(())
609/// ```
610///
611/// # Errors
612///
613/// Returns `Error::Validation` for an invalid rank or matrix shape,
614/// `Error::Extension` for an unsupported dtype or singular numerical result,
615/// and `Error::Internal` if the extension returns an unexpected number of
616/// outputs.
617pub fn full_piv_lu(
618    a: &EagerTensor,
619) -> Result<(
620    EagerTensor,
621    EagerTensor,
622    EagerTensor,
623    EagerTensor,
624    EagerTensor,
625)> {
626    let mut outputs = apply_linalg_eager(LinalgOp::FullPivLu, &[a])?.into_iter();
627    match (
628        outputs.next(),
629        outputs.next(),
630        outputs.next(),
631        outputs.next(),
632        outputs.next(),
633        outputs.next(),
634    ) {
635        (Some(p), Some(l), Some(u), Some(q), Some(parity), None) => Ok((p, l, u, q, parity)),
636        _ => Err(Error::Internal(
637            "full_piv_lu eager op returned an unexpected number of outputs".to_string(),
638        )),
639    }
640}
641
642/// Solve a linear system using complete-pivot LU behavior.
643///
644/// # Examples
645///
646/// ```rust
647/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
648/// use tenferro_linalg::EagerTensorLinalgExt;
649///
650/// let ctx = EagerRuntime::new()?;
651/// let a = EagerTensor::from_tensor_in(
652///     Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f64, 2.0, 1.0, 3.0]).unwrap(),
653///     ctx.clone(),
654/// ).unwrap();
655/// let b = EagerTensor::from_tensor_in(
656///     Tensor::from_vec_col_major(vec![2, 1], vec![-1.0_f64, 5.0]).unwrap(),
657///     ctx,
658/// ).unwrap();
659/// let x = a.full_piv_lu_solve(&b)?;
660/// assert_eq!(x.shape(), &[2, 1]);
661/// # Ok::<(), tenferro_ad::Error>(())
662/// ```
663///
664/// # Errors
665///
666/// Returns `Error::Validation` when `a` and `b` have incompatible matrix or
667/// batch shapes, `Error::Extension` for an unsupported dtype or singular
668/// system, and `Error::RuntimeState` when the backend is unavailable.
669pub fn full_piv_lu_solve(a: &EagerTensor, b: &EagerTensor) -> Result<EagerTensor> {
670    one_output(
671        apply_linalg_eager(LinalgOp::FullPivLuSolve { transpose_a: false }, &[a, b])?,
672        "full_piv_lu_solve",
673    )
674}
675
676/// Solve a linear system for eager tensors.
677///
678/// # Examples
679///
680/// ```rust
681/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
682/// use tenferro_linalg::EagerTensorLinalgExt;
683///
684/// let ctx = EagerRuntime::new()?;
685/// let a = EagerTensor::from_tensor_in(
686///     Tensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 0.0, 4.0]).unwrap(),
687///     ctx.clone(),
688/// ).unwrap();
689/// let b = EagerTensor::from_tensor_in(
690///     Tensor::from_vec_col_major(vec![2, 1], vec![4.0_f64, 8.0]).unwrap(),
691///     ctx,
692/// ).unwrap();
693/// let x = a.solve(&b)?;
694/// assert_eq!(x.shape(), &[2, 1]);
695/// # Ok::<(), tenferro_ad::Error>(())
696/// ```
697///
698/// # Errors
699///
700/// Returns `Error::Validation` for incompatible matrix, batch, or dtype
701/// metadata, `Error::Extension` for an unsupported dtype or singular system,
702/// and `Error::RuntimeState` when the backend is unavailable.
703pub fn solve(a: &EagerTensor, b: &EagerTensor) -> Result<EagerTensor> {
704    let mut factor_outputs = apply_linalg_eager(LinalgOp::LuFactor, &[a])?.into_iter();
705    let (packed_lu, pivots) = match (
706        factor_outputs.next(),
707        factor_outputs.next(),
708        factor_outputs.next(),
709        factor_outputs.next(),
710    ) {
711        (Some(packed_lu), Some(pivots), Some(_parity), None) => (packed_lu, pivots),
712        _ => {
713            return Err(Error::Internal(
714                "lu_factor eager op returned an unexpected number of outputs".to_string(),
715            ));
716        }
717    };
718    one_output(
719        apply_linalg_eager(
720            LinalgOp::LuSolvePrepared {
721                transpose_a: false,
722                conjugate_a: false,
723            },
724            &[a, &packed_lu, &pivots, b],
725        )?,
726        "solve",
727    )
728}
729
730/// Cholesky factorization for eager tensors.
731///
732/// # Examples
733///
734/// ```rust
735/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
736/// use tenferro_linalg::EagerTensorLinalgExt;
737///
738/// let ctx = EagerRuntime::new()?;
739/// let a = EagerTensor::from_tensor_in(
740///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap(),
741///     ctx,
742/// ).unwrap();
743/// let l = a.cholesky()?;
744/// assert_eq!(l.shape(), &[2, 2]);
745/// # Ok::<(), tenferro_ad::Error>(())
746/// ```
747///
748/// # Errors
749///
750/// Returns `Error::Validation` for a non-square or invalid-rank input,
751/// `Error::Extension` for an unsupported dtype or a non-positive-definite
752/// matrix, and `Error::RuntimeState` when the backend is unavailable.
753pub fn cholesky(a: &EagerTensor) -> Result<EagerTensor> {
754    one_output(apply_linalg_eager(LinalgOp::Cholesky, &[a])?, "cholesky")
755}
756
757/// Hermitian eigenvalue decomposition for eager tensors.
758///
759/// # Examples
760///
761/// ```rust
762/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
763/// use tenferro_linalg::EagerTensorLinalgExt;
764///
765/// let ctx = EagerRuntime::new()?;
766/// let a = EagerTensor::from_tensor_in(
767///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 3.0]).unwrap(),
768///     ctx,
769/// ).unwrap();
770/// let (values, vectors) = a.eigh()?;
771/// assert_eq!(values.shape(), &[2]);
772/// assert_eq!(vectors.shape(), &[2, 2]);
773/// # Ok::<(), tenferro_ad::Error>(())
774/// ```
775///
776/// # Errors
777///
778/// Returns `Error::Validation` for a non-square or invalid-rank input,
779/// `Error::Extension` for an unsupported dtype or numerical non-convergence,
780/// and `Error::RuntimeState` when the backend is unavailable.
781pub fn eigh(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor)> {
782    eigh_with_options(a, EighOptions::default())
783}
784
785/// Hermitian eigenvalue decomposition for eager tensors with explicit options.
786///
787/// `derivative_eps` regularizes derivative formulas for repeated or nearly
788/// repeated eigenvalues. It is not a backend eigensolver tolerance.
789///
790/// # Examples
791///
792/// ```rust
793/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
794/// use tenferro_linalg::{EagerTensorLinalgExt, EighGauge, EighOptions};
795///
796/// let ctx = EagerRuntime::new()?;
797/// let a = EagerTensor::from_tensor_in(
798///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 3.0]).unwrap(),
799///     ctx,
800/// ).unwrap();
801/// let (values, vectors) = a
802///     .eigh_with_options(
803///         EighOptions::default()
804///             .gauge(EighGauge::CanonicalPivot)
805///             .derivative_eps(1.0e-10),
806///     )?;
807/// assert_eq!(values.shape(), &[2]);
808/// assert_eq!(vectors.shape(), &[2, 2]);
809/// # Ok::<(), tenferro_ad::Error>(())
810/// ```
811///
812/// # Errors
813///
814/// Returns `Error::Validation` for an invalid rank, shape, or
815/// `derivative_eps`, `Error::Extension` for unsupported dtypes or numerical
816/// non-convergence, and `Error::Internal` for an output-count violation.
817pub fn eigh_with_options(
818    a: &EagerTensor,
819    options: EighOptions,
820) -> Result<(EagerTensor, EagerTensor)> {
821    validate_derivative_eps("eigh_with_options", options.derivative_eps)?;
822    two_outputs(
823        apply_linalg_eager(
824            LinalgOp::Eigh {
825                derivative_eps: options.derivative_eps,
826                gauge: options.gauge,
827            },
828            &[a],
829        )?,
830        "eigh",
831    )
832}
833
834/// General eigenvalue decomposition for eager tensors.
835///
836/// # Examples
837///
838/// ```rust
839/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
840/// use tenferro_linalg::EagerTensorLinalgExt;
841///
842/// let ctx = EagerRuntime::new()?;
843/// let a = EagerTensor::from_tensor_in(
844///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 3.0]).unwrap(),
845///     ctx,
846/// ).unwrap();
847/// let (values, vectors) = a.eig()?;
848/// assert_eq!(values.shape(), &[2]);
849/// assert_eq!(vectors.shape(), &[2, 2]);
850/// # Ok::<(), tenferro_ad::Error>(())
851/// ```
852///
853/// # Errors
854///
855/// Returns `Error::Validation` for a non-square or invalid-rank input,
856/// `Error::Extension` for an unsupported dtype or numerical non-convergence,
857/// and `Error::RuntimeState` when the backend is unavailable.
858pub fn eig(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor)> {
859    two_outputs(
860        apply_linalg_eager(
861            LinalgOp::Eig {
862                input_dtype: a.dtype(),
863            },
864            &[a],
865        )?,
866        "eig",
867    )
868}
869
870/// Triangular solve for eager tensors.
871///
872/// # Examples
873///
874/// ```rust
875/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
876/// use tenferro_linalg::EagerTensorLinalgExt;
877///
878/// let ctx = EagerRuntime::new()?;
879/// let a = EagerTensor::from_tensor_in(
880///     Tensor::from_vec_col_major(vec![2, 2], vec![2.0_f64, 0.0, 1.0, 3.0]).unwrap(),
881///     ctx.clone(),
882/// ).unwrap();
883/// let b = EagerTensor::from_tensor_in(
884///     Tensor::from_vec_col_major(vec![2, 1], vec![2.0_f64, 7.0]).unwrap(),
885///     ctx,
886/// ).unwrap();
887/// let x = a.triangular_solve(&b, true, true, false, false)?;
888/// assert_eq!(x.shape(), &[2, 1]);
889/// # Ok::<(), tenferro_ad::Error>(())
890/// ```
891///
892/// # Errors
893///
894/// Returns `Error::Validation` for incompatible matrix, batch, or dtype
895/// metadata, `Error::Extension` for an unsupported dtype or singular system,
896/// and `Error::RuntimeState` when the backend is unavailable.
897pub fn triangular_solve(
898    a: &EagerTensor,
899    b: &EagerTensor,
900    left_side: bool,
901    lower: bool,
902    transpose_a: bool,
903    unit_diagonal: bool,
904) -> Result<EagerTensor> {
905    one_output(
906        apply_linalg_eager(
907            LinalgOp::TriangularSolve {
908                left_side,
909                lower,
910                transpose_a,
911                unit_diagonal,
912            },
913            &[a, b],
914        )?,
915        "triangular_solve",
916    )
917}
918
919fn one_output(outputs: Vec<EagerTensor>, name: &str) -> Result<EagerTensor> {
920    let mut outputs = outputs.into_iter();
921    match (outputs.next(), outputs.next()) {
922        (Some(output), None) => Ok(output),
923        _ => Err(Error::Internal(format!(
924            "{name} eager op returned an unexpected number of outputs"
925        ))),
926    }
927}
928
929fn two_outputs(outputs: Vec<EagerTensor>, name: &str) -> Result<(EagerTensor, EagerTensor)> {
930    let mut outputs = outputs.into_iter();
931    match (outputs.next(), outputs.next(), outputs.next()) {
932        (Some(lhs), Some(rhs), None) => Ok((lhs, rhs)),
933        _ => Err(Error::Internal(format!(
934            "{name} eager op returned an unexpected number of outputs"
935        ))),
936    }
937}
tenferro_linalg/eager_ext.rs

tenferro_linalg/
eager_ext.rs
