1use anyhow::{anyhow, Result};
7use num_complex::{Complex32, Complex64, ComplexFloat};
8use tenferro::{DType, Tensor, TensorScalar, TensorSessionOpsExt, TypedTensor};
9use tenferro_linalg::TensorLinalgExt;
10
11use crate::context::with_default_session;
12use crate::matrix::Matrix;
13
14#[derive(Debug)]
31pub struct SvdResult<T: TensorScalar> {
32 u: TypedTensor<T>,
33 s: TypedTensor<T::Real>,
34 vt: TypedTensor<T>,
35}
36
37#[derive(Debug)]
42pub struct FullPivLuResult<T: TensorScalar> {
43 p: TypedTensor<T>,
44 l: TypedTensor<T>,
45 u: TypedTensor<T>,
46 q: TypedTensor<T>,
47}
48
49fn clone_linalg_tensor<T: TensorScalar>(tensor: &TypedTensor<T>) -> TypedTensor<T> {
50 crate::require_invariant(tensor.duplicate(), "linalg result duplication failed")
51}
52
53impl<T: TensorScalar> SvdResult<T> {
54 pub fn u(&self) -> &TypedTensor<T> {
56 &self.u
57 }
58
59 pub fn s(&self) -> &TypedTensor<T::Real> {
61 &self.s
62 }
63
64 pub fn vt(&self) -> &TypedTensor<T> {
66 &self.vt
67 }
68
69 pub fn into_parts(self) -> (TypedTensor<T>, TypedTensor<T::Real>, TypedTensor<T>) {
71 (self.u, self.s, self.vt)
72 }
73}
74
75impl<T: TensorScalar> FullPivLuResult<T> {
76 pub fn p(&self) -> &TypedTensor<T> {
78 &self.p
79 }
80
81 pub fn l(&self) -> &TypedTensor<T> {
83 &self.l
84 }
85
86 pub fn u(&self) -> &TypedTensor<T> {
88 &self.u
89 }
90
91 pub fn q(&self) -> &TypedTensor<T> {
93 &self.q
94 }
95
96 pub fn into_parts(
98 self,
99 ) -> (
100 TypedTensor<T>,
101 TypedTensor<T>,
102 TypedTensor<T>,
103 TypedTensor<T>,
104 ) {
105 (self.p, self.l, self.u, self.q)
106 }
107}
108
109impl<T: TensorScalar> Clone for SvdResult<T> {
110 fn clone(&self) -> Self {
111 Self {
112 u: clone_linalg_tensor(&self.u),
113 s: clone_linalg_tensor(&self.s),
114 vt: clone_linalg_tensor(&self.vt),
115 }
116 }
117}
118
119impl<T: TensorScalar> Clone for FullPivLuResult<T> {
120 fn clone(&self) -> Self {
121 Self {
122 p: clone_linalg_tensor(&self.p),
123 l: clone_linalg_tensor(&self.l),
124 u: clone_linalg_tensor(&self.u),
125 q: clone_linalg_tensor(&self.q),
126 }
127 }
128}
129
130#[derive(Debug, Clone)]
143pub struct FullPivLuMatrixResult<T> {
144 pub p: Matrix<T>,
146 pub l: Matrix<T>,
148 pub u: Matrix<T>,
150 pub q: Matrix<T>,
152}
153
154pub trait BackendLinalgScalar: TensorScalar {}
156
157pub trait FullPivLuScalar: BackendLinalgScalar {
172 fn solve_right_full_piv_lu(
181 lhs_values: &[Self],
182 lhs_rows: usize,
183 lhs_cols: usize,
184 pivot_values: &[Self],
185 pivot_rows: usize,
186 pivot_cols: usize,
187 ) -> std::result::Result<Vec<Self>, BackendLinalgError>;
188}
189
190macro_rules! impl_full_piv_lu_scalar {
191 ($t:ty) => {
192 impl FullPivLuScalar for $t {
193 fn solve_right_full_piv_lu(
194 lhs_values: &[Self],
195 lhs_rows: usize,
196 lhs_cols: usize,
197 pivot_values: &[Self],
198 pivot_rows: usize,
199 pivot_cols: usize,
200 ) -> std::result::Result<Vec<Self>, BackendLinalgError> {
201 if pivot_rows != pivot_cols {
202 return Err(BackendLinalgError::from(anyhow::anyhow!(
203 "full-pivot solve requires a square pivot matrix, got {}x{}",
204 pivot_rows,
205 pivot_cols
206 )));
207 }
208 if lhs_cols != pivot_rows {
209 return Err(BackendLinalgError::from(anyhow::anyhow!(
210 "cannot solve T * P = Pi1 with Pi1 shape {}x{} and P shape {}x{}",
211 lhs_rows,
212 lhs_cols,
213 pivot_rows,
214 pivot_cols
215 )));
216 }
217
218 let lhs_t = transpose_column_major(lhs_values, lhs_rows, lhs_cols);
219 let pivot_t = transpose_column_major(pivot_values, pivot_rows, pivot_cols);
220 let pivot_tensor = tenferro_tensor::Tensor::from_vec_col_major(
221 vec![pivot_cols, pivot_rows],
222 pivot_t,
223 )
224 .map_err(|e| BackendLinalgError::from(anyhow::Error::new(e)))?;
225 let lhs_tensor =
226 tenferro_tensor::Tensor::from_vec_col_major(vec![lhs_cols, lhs_rows], lhs_t)
227 .map_err(|e| BackendLinalgError::from(anyhow::Error::new(e)))?;
228 let solved_t = with_default_session(|session| {
229 pivot_tensor.full_piv_lu_solve(&lhs_tensor, session)
230 })
231 .map_err(|e| {
232 BackendLinalgError::from(anyhow::anyhow!("full_piv_lu_solve failed: {e}"))
233 })?;
234
235 let solved_t_values = solved_t.as_slice::<Self>().map_err(|e| {
236 BackendLinalgError::from(anyhow::anyhow!(
237 "full_piv_lu_solve returned unexpected dtype: {e}"
238 ))
239 })?;
240 Ok(transpose_column_major(solved_t_values, lhs_cols, lhs_rows))
241 }
242 }
243 };
244}
245
246impl_full_piv_lu_scalar!(f32);
247impl_full_piv_lu_scalar!(f64);
248impl_full_piv_lu_scalar!(num_complex::Complex32);
249impl_full_piv_lu_scalar!(num_complex::Complex64);
250
251fn transpose_column_major<T: Copy + num_traits::Zero>(
253 values: &[T],
254 nrows: usize,
255 ncols: usize,
256) -> Vec<T> {
257 let mut out = vec![T::zero(); nrows * ncols];
258 for col in 0..ncols {
259 for row in 0..nrows {
260 out[col + ncols * row] = values[row + nrows * col];
261 }
262 }
263 out
264}
265
266impl<T: TensorScalar> BackendLinalgScalar for T {}
267
268pub trait MatrixSolveScalar: BackendLinalgScalar + crate::matrix::MatrixScalar {
282 #[doc(hidden)]
283 fn solve_matrix_impl(a: &Matrix<Self>, b: &Matrix<Self>) -> Result<Matrix<Self>>;
284
285 #[doc(hidden)]
286 fn solve_matrix_owned_impl(a: Matrix<Self>, b: Matrix<Self>) -> Result<Matrix<Self>> {
287 Self::solve_matrix_impl(&a, &b)
288 }
289}
290
291#[derive(Debug, thiserror::Error)]
303#[error("backend linear algebra failed: {source}")]
304pub struct BackendLinalgError {
305 #[source]
307 pub source: anyhow::Error,
308}
309
310impl From<anyhow::Error> for BackendLinalgError {
311 fn from(source: anyhow::Error) -> Self {
312 Self { source }
313 }
314}
315
316pub trait MatrixTriangularSolveScalar: BackendLinalgScalar + crate::matrix::MatrixScalar {
331 #[doc(hidden)]
332 fn triangular_solve_matrix_impl(
333 a: &Matrix<Self>,
334 b: &Matrix<Self>,
335 left_side: bool,
336 lower: bool,
337 transpose_a: bool,
338 unit_diagonal: bool,
339 ) -> Result<Matrix<Self>>;
340
341 #[doc(hidden)]
342 fn triangular_solve_matrix_owned_impl(
343 a: Matrix<Self>,
344 b: Matrix<Self>,
345 left_side: bool,
346 lower: bool,
347 transpose_a: bool,
348 unit_diagonal: bool,
349 ) -> Result<Matrix<Self>> {
350 Self::triangular_solve_matrix_impl(&a, &b, left_side, lower, transpose_a, unit_diagonal)
351 }
352}
353
354#[derive(Debug, Clone, Copy, PartialEq)]
372pub struct SrcErrorEstimate {
373 pub error: f64,
375 pub norm: f64,
377}
378
379pub fn src_error_estimate<T>(
411 r: &Matrix<T>,
412) -> std::result::Result<SrcErrorEstimate, BackendLinalgError>
413where
414 T: MatrixTriangularSolveScalar + ComplexFloat,
415{
416 let inverse_adjoint = src_inverse_adjoint(r)?;
417 src_error_estimate_from_inverse_adjoint(r, &inverse_adjoint)
418}
419
420pub fn src_error_estimate_general<T>(
442 factor: &Matrix<T>,
443) -> std::result::Result<SrcErrorEstimate, BackendLinalgError>
444where
445 T: MatrixSolveScalar + ComplexFloat,
446{
447 let nrows = factor.nrows();
448 let ncols = factor.ncols();
449 if nrows != ncols {
450 return Err(BackendLinalgError::from(anyhow!(
451 "SRC estimator requires a square factor, got {nrows}x{ncols}"
452 )));
453 }
454 if nrows == 0 {
455 return Err(BackendLinalgError::from(anyhow!(
456 "SRC estimator requires a non-empty factor"
457 )));
458 }
459 if factor
460 .as_col_major_slice()
461 .iter()
462 .any(|value| !value.matrix_abs_sq().is_finite())
463 {
464 return Err(BackendLinalgError::from(anyhow!(
465 "SRC estimator requires finite entries in its factor"
466 )));
467 }
468 let mut adjoint = Matrix::zeros(nrows, ncols);
469 let mut identity = Matrix::zeros(nrows, ncols);
470 for col in 0..ncols {
471 for row in 0..nrows {
472 adjoint[[row, col]] = factor[[col, row]].conj();
473 }
474 identity[[col, col]] = T::one();
475 }
476 let inverse_adjoint = solve_matrix(&adjoint, &identity).map_err(|error| {
477 BackendLinalgError::from(anyhow!("SRC inverse-adjoint general solve failed: {error}"))
478 })?;
479 src_error_estimate_from_inverse_adjoint(factor, &inverse_adjoint)
480}
481
482pub(crate) fn src_inverse_adjoint<T>(
489 r: &Matrix<T>,
490) -> std::result::Result<Matrix<T>, BackendLinalgError>
491where
492 T: MatrixTriangularSolveScalar + ComplexFloat,
493{
494 let nrows = r.nrows();
495 let ncols = r.ncols();
496 if nrows != ncols {
497 return Err(BackendLinalgError::from(anyhow!(
498 "SRC estimator requires a square R, got {nrows}x{ncols}"
499 )));
500 }
501 if nrows == 0 {
502 return Err(BackendLinalgError::from(anyhow!(
503 "SRC estimator requires a non-empty R"
504 )));
505 }
506
507 for col in 0..ncols {
514 let diagonal_sq = r[[col, col]].matrix_abs_sq();
515 if !diagonal_sq.is_finite() || diagonal_sq == 0.0 {
516 return Err(BackendLinalgError::from(anyhow!(
517 "SRC estimator requires a finite, nonzero diagonal in R at ({col}, {col})"
518 )));
519 }
520 for row in col + 1..nrows {
521 if r[[row, col]].matrix_abs_sq() != 0.0 {
522 return Err(BackendLinalgError::from(anyhow!(
523 "SRC triangular estimator requires upper-triangular R; entry ({row}, {col}) is nonzero"
524 )));
525 }
526 }
527 }
528
529 let mut adjoint = Matrix::zeros(nrows, ncols);
530 for col in 0..ncols {
531 for row in 0..nrows {
532 adjoint[[row, col]] = r[[col, row]].conj();
533 }
534 }
535 let mut identity = Matrix::zeros(nrows, ncols);
536 for diagonal in 0..nrows {
537 identity[[diagonal, diagonal]] = T::one();
538 }
539
540 let inverse_adjoint = triangular_solve_matrix(&adjoint, &identity, true, true, false, false)
541 .map_err(|error| {
542 BackendLinalgError::from(anyhow!(
543 "SRC inverse-adjoint triangular solve failed: {error}"
544 ))
545 })?;
546 Ok(inverse_adjoint)
547}
548
549pub(crate) fn src_error_estimate_from_inverse_adjoint<T>(
555 r: &Matrix<T>,
556 inverse_adjoint: &Matrix<T>,
557) -> std::result::Result<SrcErrorEstimate, BackendLinalgError>
558where
559 T: crate::matrix::MatrixScalar + ComplexFloat,
560{
561 let nrows = r.nrows();
562 let ncols = r.ncols();
563 if nrows != ncols || inverse_adjoint.nrows() != nrows || inverse_adjoint.ncols() != ncols {
564 return Err(BackendLinalgError::from(anyhow!(
565 "SRC estimator requires matching square R and inverse-adjoint factors"
566 )));
567 }
568 if nrows == 0 {
569 return Err(BackendLinalgError::from(anyhow!(
570 "SRC estimator requires a non-empty R"
571 )));
572 }
573
574 let mut norm_sq = 0.0_f64;
575 for value in r.as_col_major_slice() {
576 norm_sq += value.matrix_abs_sq();
577 }
578 if !norm_sq.is_finite() {
579 return Err(BackendLinalgError::from(anyhow!(
580 "SRC estimator requires finite entries in R"
581 )));
582 }
583 let mut inverse_column_error_sq = 0.0_f64;
584 for col in 0..ncols {
585 let column_norm_sq = (0..nrows)
586 .map(|row| inverse_adjoint[[row, col]].matrix_abs_sq())
587 .sum::<f64>();
588 if !column_norm_sq.is_finite() || column_norm_sq == 0.0 {
589 return Err(BackendLinalgError::from(anyhow!(
590 "SRC inverse-adjoint solve returned an invalid column norm at column {col}"
591 )));
592 }
593 inverse_column_error_sq += 1.0 / column_norm_sq;
594 }
595
596 let sketch_width = ncols as f64;
597 let error_sq = inverse_column_error_sq / sketch_width;
598 let norm_estimate_sq = norm_sq / sketch_width;
599 if !error_sq.is_finite() || !norm_estimate_sq.is_finite() {
600 return Err(BackendLinalgError::from(anyhow!(
601 "SRC estimator produced a non-finite estimate"
602 )));
603 }
604 Ok(SrcErrorEstimate {
605 error: error_sq.sqrt(),
606 norm: norm_estimate_sq.sqrt(),
607 })
608}
609
610fn solve_matrix_direct<T>(a: &Matrix<T>, b: &Matrix<T>) -> Result<Matrix<T>>
611where
612 T: BackendLinalgScalar + Copy,
613 Tensor: From<TypedTensor<T>>,
614{
615 solve_matrix_direct_owned(a.clone(), b.clone())
616}
617
618fn solve_matrix_direct_owned<T>(a: Matrix<T>, b: Matrix<T>) -> Result<Matrix<T>>
619where
620 T: BackendLinalgScalar + Copy,
621 Tensor: From<TypedTensor<T>>,
622{
623 let a_tensor: Tensor = a.into_typed_tensor().into();
624 let b_tensor: Tensor = b.into_typed_tensor().into();
625 let result = with_default_session(|session| a_tensor.solve(&b_tensor, session))
626 .map_err(|e| anyhow!("linear solve failed via tenferro-tensor: {e}"))?;
627 let x = try_into_typed_result::<T>("solve", result)?;
628 typed_tensor_to_matrix("solve", x)
629}
630
631fn triangular_solve_matrix_direct<T>(
632 a: &Matrix<T>,
633 b: &Matrix<T>,
634 left_side: bool,
635 lower: bool,
636 transpose_a: bool,
637 unit_diagonal: bool,
638) -> Result<Matrix<T>>
639where
640 T: BackendLinalgScalar + Copy,
641 Tensor: From<TypedTensor<T>>,
642{
643 triangular_solve_matrix_direct_owned(
644 a.clone(),
645 b.clone(),
646 left_side,
647 lower,
648 transpose_a,
649 unit_diagonal,
650 )
651}
652
653fn triangular_solve_matrix_direct_owned<T>(
654 a: Matrix<T>,
655 b: Matrix<T>,
656 left_side: bool,
657 lower: bool,
658 transpose_a: bool,
659 unit_diagonal: bool,
660) -> Result<Matrix<T>>
661where
662 T: BackendLinalgScalar + Copy,
663 Tensor: From<TypedTensor<T>>,
664{
665 let a_tensor: Tensor = a.into_typed_tensor().into();
666 let b_tensor: Tensor = b.into_typed_tensor().into();
667 let result = with_default_session(|session| {
668 a_tensor.triangular_solve(
669 &b_tensor,
670 left_side,
671 lower,
672 transpose_a,
673 unit_diagonal,
674 session,
675 )
676 })
677 .map_err(|e| anyhow!("triangular solve failed via tenferro-tensor: {e}"))?;
678 let x = try_into_typed_result::<T>("triangular_solve", result)?;
679 typed_tensor_to_matrix("triangular_solve", x)
680}
681
682impl MatrixSolveScalar for f64 {
683 fn solve_matrix_impl(a: &Matrix<Self>, b: &Matrix<Self>) -> Result<Matrix<Self>> {
684 solve_matrix_direct(a, b)
685 }
686
687 fn solve_matrix_owned_impl(a: Matrix<Self>, b: Matrix<Self>) -> Result<Matrix<Self>> {
688 solve_matrix_direct_owned(a, b)
689 }
690}
691
692impl MatrixTriangularSolveScalar for f64 {
693 fn triangular_solve_matrix_impl(
694 a: &Matrix<Self>,
695 b: &Matrix<Self>,
696 left_side: bool,
697 lower: bool,
698 transpose_a: bool,
699 unit_diagonal: bool,
700 ) -> Result<Matrix<Self>> {
701 triangular_solve_matrix_direct(a, b, left_side, lower, transpose_a, unit_diagonal)
702 }
703
704 fn triangular_solve_matrix_owned_impl(
705 a: Matrix<Self>,
706 b: Matrix<Self>,
707 left_side: bool,
708 lower: bool,
709 transpose_a: bool,
710 unit_diagonal: bool,
711 ) -> Result<Matrix<Self>> {
712 triangular_solve_matrix_direct_owned(a, b, left_side, lower, transpose_a, unit_diagonal)
713 }
714}
715
716impl MatrixSolveScalar for Complex64 {
717 fn solve_matrix_impl(a: &Matrix<Self>, b: &Matrix<Self>) -> Result<Matrix<Self>> {
718 solve_matrix_direct(a, b)
719 }
720
721 fn solve_matrix_owned_impl(a: Matrix<Self>, b: Matrix<Self>) -> Result<Matrix<Self>> {
722 solve_matrix_direct_owned(a, b)
723 }
724}
725
726impl MatrixTriangularSolveScalar for Complex64 {
727 fn triangular_solve_matrix_impl(
728 a: &Matrix<Self>,
729 b: &Matrix<Self>,
730 left_side: bool,
731 lower: bool,
732 transpose_a: bool,
733 unit_diagonal: bool,
734 ) -> Result<Matrix<Self>> {
735 triangular_solve_matrix_direct(a, b, left_side, lower, transpose_a, unit_diagonal)
736 }
737
738 fn triangular_solve_matrix_owned_impl(
739 a: Matrix<Self>,
740 b: Matrix<Self>,
741 left_side: bool,
742 lower: bool,
743 transpose_a: bool,
744 unit_diagonal: bool,
745 ) -> Result<Matrix<Self>> {
746 triangular_solve_matrix_direct_owned(a, b, left_side, lower, transpose_a, unit_diagonal)
747 }
748}
749
750impl MatrixSolveScalar for f32 {
751 fn solve_matrix_impl(a: &Matrix<Self>, b: &Matrix<Self>) -> Result<Matrix<Self>> {
752 let a64 = Matrix::from_col_major_vec(
753 a.nrows(),
754 a.ncols(),
755 a.as_col_major_slice()
756 .iter()
757 .map(|&value| value as f64)
758 .collect(),
759 );
760 let b64 = Matrix::from_col_major_vec(
761 b.nrows(),
762 b.ncols(),
763 b.as_col_major_slice()
764 .iter()
765 .map(|&value| value as f64)
766 .collect(),
767 );
768 let x64 = solve_matrix_direct(&a64, &b64)?;
769 Ok(Matrix::from_col_major_vec(
770 x64.nrows(),
771 x64.ncols(),
772 x64.as_col_major_slice()
773 .iter()
774 .map(|&value| value as f32)
775 .collect(),
776 ))
777 }
778}
779
780impl MatrixTriangularSolveScalar for f32 {
781 fn triangular_solve_matrix_impl(
782 a: &Matrix<Self>,
783 b: &Matrix<Self>,
784 left_side: bool,
785 lower: bool,
786 transpose_a: bool,
787 unit_diagonal: bool,
788 ) -> Result<Matrix<Self>> {
789 let a64 = Matrix::from_col_major_vec(
790 a.nrows(),
791 a.ncols(),
792 a.as_col_major_slice()
793 .iter()
794 .map(|&value| value as f64)
795 .collect(),
796 );
797 let b64 = Matrix::from_col_major_vec(
798 b.nrows(),
799 b.ncols(),
800 b.as_col_major_slice()
801 .iter()
802 .map(|&value| value as f64)
803 .collect(),
804 );
805 let x64 = triangular_solve_matrix_direct(
806 &a64,
807 &b64,
808 left_side,
809 lower,
810 transpose_a,
811 unit_diagonal,
812 )?;
813 Ok(Matrix::from_col_major_vec(
814 x64.nrows(),
815 x64.ncols(),
816 x64.as_col_major_slice()
817 .iter()
818 .map(|&value| value as f32)
819 .collect(),
820 ))
821 }
822}
823
824impl MatrixSolveScalar for Complex32 {
825 fn solve_matrix_impl(a: &Matrix<Self>, b: &Matrix<Self>) -> Result<Matrix<Self>> {
826 let a64 = Matrix::from_col_major_vec(
827 a.nrows(),
828 a.ncols(),
829 a.as_col_major_slice()
830 .iter()
831 .map(|&value| Complex64::new(value.re as f64, value.im as f64))
832 .collect(),
833 );
834 let b64 = Matrix::from_col_major_vec(
835 b.nrows(),
836 b.ncols(),
837 b.as_col_major_slice()
838 .iter()
839 .map(|&value| Complex64::new(value.re as f64, value.im as f64))
840 .collect(),
841 );
842 let x64 = solve_matrix_direct(&a64, &b64)?;
843 Ok(Matrix::from_col_major_vec(
844 x64.nrows(),
845 x64.ncols(),
846 x64.as_col_major_slice()
847 .iter()
848 .map(|&value| Complex32::new(value.re as f32, value.im as f32))
849 .collect(),
850 ))
851 }
852}
853
854impl MatrixTriangularSolveScalar for Complex32 {
855 fn triangular_solve_matrix_impl(
856 a: &Matrix<Self>,
857 b: &Matrix<Self>,
858 left_side: bool,
859 lower: bool,
860 transpose_a: bool,
861 unit_diagonal: bool,
862 ) -> Result<Matrix<Self>> {
863 let a64 = Matrix::from_col_major_vec(
864 a.nrows(),
865 a.ncols(),
866 a.as_col_major_slice()
867 .iter()
868 .map(|&value| Complex64::new(value.re as f64, value.im as f64))
869 .collect(),
870 );
871 let b64 = Matrix::from_col_major_vec(
872 b.nrows(),
873 b.ncols(),
874 b.as_col_major_slice()
875 .iter()
876 .map(|&value| Complex64::new(value.re as f64, value.im as f64))
877 .collect(),
878 );
879 let x64 = triangular_solve_matrix_direct(
880 &a64,
881 &b64,
882 left_side,
883 lower,
884 transpose_a,
885 unit_diagonal,
886 )?;
887 Ok(Matrix::from_col_major_vec(
888 x64.nrows(),
889 x64.ncols(),
890 x64.as_col_major_slice()
891 .iter()
892 .map(|&value| Complex32::new(value.re as f32, value.im as f32))
893 .collect(),
894 ))
895 }
896}
897
898fn tensor_scalar_dtype<T: TensorScalar>() -> DType {
899 T::dtype()
900}
901
902fn try_into_typed_result<T: TensorScalar>(
903 op: &'static str,
904 tensor: Tensor,
905) -> Result<TypedTensor<T>> {
906 let actual = tensor.dtype();
907 T::into_typed(tensor).map_err(|source| {
908 anyhow!(
909 "{op}: dtype mismatch lhs={actual:?} rhs={:?}: {source}",
910 tensor_scalar_dtype::<T>()
911 )
912 })
913}
914
915fn convert_for_typed<T: TensorScalar>(op: &'static str, tensor: Tensor) -> Result<TypedTensor<T>> {
916 let expected = tensor_scalar_dtype::<T>();
917 let tensor = if tensor.dtype() == expected {
918 tensor
919 } else {
920 with_default_session(|session| tensor.convert(expected, session))
921 .map_err(|e| anyhow!("{op}: dtype conversion to {expected:?} failed: {e}"))?
922 };
923 try_into_typed_result::<T>(op, tensor)
924}
925
926fn matrix_to_typed_tensor<T>(matrix: &Matrix<T>) -> TypedTensor<T>
927where
928 T: TensorScalar + Copy,
929{
930 crate::require_invariant(
931 TypedTensor::from_vec_col_major(
932 vec![matrix.nrows(), matrix.ncols()],
933 matrix.as_col_major_slice().to_vec(),
934 ),
935 "validated matrix rejected by tenferro",
936 )
937}
938
939fn typed_tensor_to_matrix<T>(op: &'static str, tensor: TypedTensor<T>) -> Result<Matrix<T>>
940where
941 T: TensorScalar + Copy,
942{
943 Matrix::try_from_typed_tensor(tensor).map_err(|err| anyhow!("{op}: {err}"))
944}
945
946fn require_host_linalg_tensor<T: TensorScalar>(
947 op: &'static str,
948 tensor: TypedTensor<T>,
949) -> Result<TypedTensor<T>> {
950 tensor
951 .host_data()
952 .map_err(|error| anyhow!("{op}: result must be host-backed: {error}"))?;
953 Ok(tensor)
954}
955
956pub fn svd_backend<T>(a: &TypedTensor<T>) -> std::result::Result<SvdResult<T>, BackendLinalgError>
963where
964 T: BackendLinalgScalar,
965{
966 let tensor = T::into_tensor(
967 a.shape().to_vec(),
968 a.host_data()
969 .map_err(|e| anyhow!("SVD input host access failed: {e}"))?
970 .to_vec(),
971 )
972 .map_err(|e| anyhow!("SVD input tensor construction failed: {e}"))?;
973 let (u, s, vt) = with_default_session(|session| tensor.svd(session))
974 .map_err(|e| anyhow!("SVD computation failed via tenferro-tensor: {e}"))?;
975 Ok(SvdResult {
976 u: require_host_linalg_tensor("svd", convert_for_typed::<T>("svd", u)?)?,
977 s: require_host_linalg_tensor("svd", convert_for_typed::<T::Real>("svd", s)?)?,
978 vt: require_host_linalg_tensor("svd", convert_for_typed::<T>("svd", vt)?)?,
979 })
980}
981
982pub fn qr_backend<T>(
990 a: TypedTensor<T>,
991) -> std::result::Result<(TypedTensor<T>, TypedTensor<T>), BackendLinalgError>
992where
993 T: BackendLinalgScalar,
994{
995 let (shape, data) = a
996 .into_vec_col_major()
997 .map_err(|e| anyhow!("QR input host access failed: {e}"))?;
998 let tensor = T::into_tensor(shape, data)
999 .map_err(|e| anyhow!("QR input tensor construction failed: {e}"))?;
1000 let (q, r) = with_default_session(|session| tensor.qr(session))
1001 .map_err(|e| anyhow!("QR computation failed via tenferro-tensor: {e}"))?;
1002 Ok((
1003 convert_for_typed::<T>("qr", q)?,
1004 convert_for_typed::<T>("qr", r)?,
1005 ))
1006}
1007
1008pub fn solve_backend<T>(
1014 a: &TypedTensor<T>,
1015 b: &TypedTensor<T>,
1016) -> std::result::Result<TypedTensor<T>, BackendLinalgError>
1017where
1018 T: BackendLinalgScalar,
1019{
1020 let a_tensor = T::into_tensor(
1021 a.shape().to_vec(),
1022 a.host_data()
1023 .map_err(|e| anyhow!("solve input host access failed: {e}"))?
1024 .to_vec(),
1025 )
1026 .map_err(|e| anyhow!("solve lhs tensor construction failed: {e}"))?;
1027 let b_tensor = T::into_tensor(
1028 b.shape().to_vec(),
1029 b.host_data()
1030 .map_err(|e| anyhow!("solve rhs host access failed: {e}"))?
1031 .to_vec(),
1032 )
1033 .map_err(|e| anyhow!("solve rhs tensor construction failed: {e}"))?;
1034 let result = with_default_session(|session| a_tensor.solve(&b_tensor, session))
1035 .map_err(|e| anyhow!("linear solve failed via tenferro-tensor: {e}"))?;
1036 try_into_typed_result::<T>("solve", result).map_err(BackendLinalgError::from)
1037}
1038
1039pub fn triangular_solve_backend<T>(
1049 a: &TypedTensor<T>,
1050 b: &TypedTensor<T>,
1051 left_side: bool,
1052 lower: bool,
1053 transpose_a: bool,
1054 unit_diagonal: bool,
1055) -> std::result::Result<TypedTensor<T>, BackendLinalgError>
1056where
1057 T: BackendLinalgScalar,
1058{
1059 let a_tensor = T::into_tensor(
1060 a.shape().to_vec(),
1061 a.host_data()
1062 .map_err(|e| anyhow!("triangular solve input host access failed: {e}"))?
1063 .to_vec(),
1064 )
1065 .map_err(|e| anyhow!("triangular solve lhs tensor construction failed: {e}"))?;
1066 let b_tensor = T::into_tensor(
1067 b.shape().to_vec(),
1068 b.host_data()
1069 .map_err(|e| anyhow!("triangular solve rhs host access failed: {e}"))?
1070 .to_vec(),
1071 )
1072 .map_err(|e| anyhow!("triangular solve rhs tensor construction failed: {e}"))?;
1073 let result = with_default_session(|session| {
1074 a_tensor.triangular_solve(
1075 &b_tensor,
1076 left_side,
1077 lower,
1078 transpose_a,
1079 unit_diagonal,
1080 session,
1081 )
1082 })
1083 .map_err(|e| anyhow!("triangular solve failed via tenferro-tensor: {e}"))?;
1084 try_into_typed_result::<T>("triangular_solve", result).map_err(BackendLinalgError::from)
1085}
1086
1087pub fn solve_matrix<T>(
1106 a: &Matrix<T>,
1107 b: &Matrix<T>,
1108) -> std::result::Result<Matrix<T>, BackendLinalgError>
1109where
1110 T: MatrixSolveScalar,
1111{
1112 T::solve_matrix_impl(a, b).map_err(BackendLinalgError::from)
1113}
1114
1115pub fn solve_matrix_owned<T>(
1135 a: Matrix<T>,
1136 b: Matrix<T>,
1137) -> std::result::Result<Matrix<T>, BackendLinalgError>
1138where
1139 T: MatrixSolveScalar,
1140{
1141 T::solve_matrix_owned_impl(a, b).map_err(BackendLinalgError::from)
1142}
1143
1144pub fn triangular_solve_matrix<T>(
1165 a: &Matrix<T>,
1166 b: &Matrix<T>,
1167 left_side: bool,
1168 lower: bool,
1169 transpose_a: bool,
1170 unit_diagonal: bool,
1171) -> std::result::Result<Matrix<T>, BackendLinalgError>
1172where
1173 T: MatrixTriangularSolveScalar,
1174{
1175 T::triangular_solve_matrix_impl(a, b, left_side, lower, transpose_a, unit_diagonal)
1176 .map_err(BackendLinalgError::from)
1177}
1178
1179pub fn triangular_solve_matrix_owned<T>(
1200 a: Matrix<T>,
1201 b: Matrix<T>,
1202 left_side: bool,
1203 lower: bool,
1204 transpose_a: bool,
1205 unit_diagonal: bool,
1206) -> std::result::Result<Matrix<T>, BackendLinalgError>
1207where
1208 T: MatrixTriangularSolveScalar,
1209{
1210 T::triangular_solve_matrix_owned_impl(a, b, left_side, lower, transpose_a, unit_diagonal)
1211 .map_err(BackendLinalgError::from)
1212}
1213
1214fn full_piv_lu_tensor<T>(
1215 tensor: Tensor,
1216) -> std::result::Result<FullPivLuResult<T>, BackendLinalgError>
1217where
1218 T: BackendLinalgScalar,
1219{
1220 let (p, l, u, q, _parity) = with_default_session(|session| tensor.full_piv_lu(session))
1221 .map_err(|e| anyhow!("complete-pivoting LU failed via tenferro-tensor: {e}"))?;
1222 Ok(FullPivLuResult {
1223 p: require_host_linalg_tensor("full_piv_lu", convert_for_typed::<T>("full_piv_lu", p)?)?,
1224 l: require_host_linalg_tensor("full_piv_lu", convert_for_typed::<T>("full_piv_lu", l)?)?,
1225 u: require_host_linalg_tensor("full_piv_lu", convert_for_typed::<T>("full_piv_lu", u)?)?,
1226 q: require_host_linalg_tensor("full_piv_lu", convert_for_typed::<T>("full_piv_lu", q)?)?,
1227 })
1228}
1229
1230pub fn full_piv_lu_backend<T>(
1237 a: &TypedTensor<T>,
1238) -> std::result::Result<FullPivLuResult<T>, BackendLinalgError>
1239where
1240 T: BackendLinalgScalar,
1241{
1242 let tensor = T::into_tensor(
1243 a.shape().to_vec(),
1244 a.host_data()
1245 .map_err(|e| anyhow!("LU input host access failed: {e}"))?
1246 .to_vec(),
1247 )
1248 .map_err(|e| anyhow!("LU input tensor construction failed: {e}"))?;
1249 full_piv_lu_tensor(tensor)
1250}
1251
1252pub fn full_piv_lu_matrix<T>(
1270 a: &Matrix<T>,
1271) -> std::result::Result<FullPivLuMatrixResult<T>, BackendLinalgError>
1272where
1273 T: BackendLinalgScalar + Copy,
1274{
1275 let tensor = matrix_to_typed_tensor(a);
1276 let decomp = full_piv_lu_backend(&tensor)?;
1277 Ok(FullPivLuMatrixResult {
1278 p: typed_tensor_to_matrix("full_piv_lu", decomp.p)?,
1279 l: typed_tensor_to_matrix("full_piv_lu", decomp.l)?,
1280 u: typed_tensor_to_matrix("full_piv_lu", decomp.u)?,
1281 q: typed_tensor_to_matrix("full_piv_lu", decomp.q)?,
1282 })
1283}
1284
1285pub fn full_piv_lu_matrix_owned<T>(
1337 a: Matrix<T>,
1338) -> std::result::Result<FullPivLuMatrixResult<T>, BackendLinalgError>
1339where
1340 T: BackendLinalgScalar + Copy,
1341{
1342 let tensor = T::typed_tensor_into_tensor(a.into_typed_tensor());
1343 let decomp = full_piv_lu_tensor(tensor)?;
1344 Ok(FullPivLuMatrixResult {
1345 p: typed_tensor_to_matrix("full_piv_lu", decomp.p)?,
1346 l: typed_tensor_to_matrix("full_piv_lu", decomp.l)?,
1347 u: typed_tensor_to_matrix("full_piv_lu", decomp.u)?,
1348 q: typed_tensor_to_matrix("full_piv_lu", decomp.q)?,
1349 })
1350}
1351
1352#[cfg(test)]
1353mod tests;