Skip to main content

tenferro_linalg/
backend.rs

1use tenferro_tensor::{Tensor, TensorBackend, TensorView};
2
3use crate::extension::{
4    apply_eigh_gauge, apply_qr_gauge, apply_svd_gauge, validate_derivative_eps, EighOptions,
5    QrOptions, SvdOptions,
6};
7
8/// Build the shared "unsupported dtype" backend-failure error used by the
9/// linalg backends.
10///
11/// Both the CPU and GPU linalg backends reject integer and boolean dtypes for
12/// floating-point decompositions with an identical message; this helper keeps
13/// that error construction in one place.
14pub(crate) fn unsupported_dtype(
15    op: &'static str,
16    dtype: tenferro_tensor::DType,
17) -> tenferro_tensor::Error {
18    tenferro_tensor::Error::backend_failure(op, format!("unsupported dtype {dtype:?}"))
19}
20
21/// Backend surface required by the linalg extension runtime.
22///
23/// # Examples
24///
25/// ```rust
26/// use tenferro_linalg::backend::LinalgBackend;
27/// use tenferro_cpu::CpuBackend;
28///
29/// fn accepts_linalg_backend<B: LinalgBackend>(_backend: &mut B) {}
30///
31/// let mut backend = CpuBackend::new();
32/// accepts_linalg_backend(&mut backend);
33/// ```
34pub trait LinalgBackend: TensorBackend {
35    /// Compute a Cholesky factorization.
36    fn cholesky(&mut self, input: &Tensor) -> tenferro_tensor::Result<Tensor>;
37
38    /// Solve a triangular linear system with explicit side, triangle,
39    /// transpose, and unit-diagonal flags.
40    fn triangular_solve(
41        &mut self,
42        a: &Tensor,
43        b: &Tensor,
44        left_side: bool,
45        lower: bool,
46        transpose_a: bool,
47        unit_diagonal: bool,
48    ) -> tenferro_tensor::Result<Tensor>;
49
50    /// Compute public LU outputs `(P, L, U, parity)`.
51    fn lu(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>>;
52
53    #[doc(hidden)]
54    fn lu_factor(&mut self, _input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>> {
55        Err(tenferro_tensor::Error::backend_failure(
56            "lu_factor",
57            format!(
58                "backend {} does not implement internal packed LU factorization",
59                std::any::type_name::<Self>()
60            ),
61        ))
62    }
63
64    /// Compute complete-pivot LU outputs `(P, L, U, Q, parity)`.
65    ///
66    /// The reconstruction convention is `A = P^T * L * U * Q`, equivalently
67    /// `P * A * Q^T = L * U`. `parity` is a scalar real tensor containing
68    /// `+1` or `-1`: `F32` for `F32`/`C32` inputs and `F64` for `F64`/`C64`
69    /// inputs.
70    fn full_piv_lu(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>>;
71
72    /// Solve a linear system through the complete-pivot LU path.
73    ///
74    /// With `transpose_a = false`, this solves `A * x = b`. With
75    /// `transpose_a = true`, this solves `A^T * x = b`.
76    fn full_piv_lu_solve(
77        &mut self,
78        a: &Tensor,
79        b: &Tensor,
80        transpose_a: bool,
81    ) -> tenferro_tensor::Result<Tensor>;
82
83    /// Compute public SVD outputs `(U, S, Vt)`.
84    fn svd(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>>;
85
86    /// Compute public SVD outputs `(U, S, Vt)` with explicit options.
87    ///
88    /// `derivative_eps` is validated for API consistency, but concrete backend
89    /// execution does not perform AD. `gauge` controls optional singular-vector
90    /// post-processing.
91    ///
92    /// # Examples
93    ///
94    /// ```rust
95    /// use tenferro_cpu::CpuBackend;
96    /// use tenferro_linalg::{LinalgBackend, SvdGauge, SvdOptions};
97    /// use tenferro_tensor::Tensor;
98    ///
99    /// let input = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0])?;
100    /// let mut backend = CpuBackend::new();
101    /// let outputs = backend.svd_with_options(
102    ///     &input,
103    ///     SvdOptions::default().gauge(SvdGauge::CanonicalPivot),
104    /// )?;
105    /// assert_eq!(outputs[1].shape(), &[2]);
106    /// # Ok::<(), tenferro_tensor::Error>(())
107    /// ```
108    fn svd_with_options(
109        &mut self,
110        input: &Tensor,
111        options: SvdOptions,
112    ) -> tenferro_tensor::Result<Vec<Tensor>> {
113        validate_derivative_eps("svd_with_options", options.derivative_eps)?;
114        let mut outputs = self.svd(input)?;
115        apply_svd_gauge(options.gauge, &mut outputs)?;
116        Ok(outputs)
117    }
118
119    #[doc(hidden)]
120    fn svd_values(&mut self, _input: &Tensor) -> tenferro_tensor::Result<Tensor> {
121        Err(tenferro_tensor::Error::backend_failure(
122            "svd_values",
123            format!(
124                "backend {} does not implement internal singular-values-only decomposition",
125                std::any::type_name::<Self>()
126            ),
127        ))
128    }
129
130    /// Compute a singular value decomposition from a borrowed tensor view.
131    ///
132    /// Backends may canonicalize the view inside the same placement family, but
133    /// must not silently transfer between CPU and GPU memory.
134    ///
135    /// # Examples
136    ///
137    /// ```rust
138    /// use tenferro_linalg::LinalgBackend;
139    /// use tenferro_cpu::CpuBackend;
140    /// use tenferro_tensor::{TensorView, TypedTensor};
141    ///
142    /// let input = TypedTensor::<f64>::from_vec_col_major(
143    ///     vec![2, 2],
144    ///     vec![1.0, 0.0, 0.0, 2.0],
145    /// )?;
146    /// let mut backend = CpuBackend::new();
147    /// let outputs = backend.svd_read(TensorView::F64(input.as_view()))?;
148    /// assert_eq!(outputs[1].shape(), &[2]);
149    /// # Ok::<(), tenferro_tensor::Error>(())
150    /// ```
151    fn svd_read(&mut self, _input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
152        Err(tenferro_tensor::Error::backend_failure(
153            "svd",
154            "backend does not accept borrowed tensor views at this execution boundary",
155        ))
156    }
157
158    /// Compute public QR outputs `(Q, R)`.
159    ///
160    /// QR is thin: for an `m x n` input, `Q` has shape `m x min(m, n)` and
161    /// `R` has shape `min(m, n) x n`.
162    fn qr(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>>;
163
164    /// Compute public QR outputs `(Q, R)` with explicit options.
165    ///
166    /// `gauge` controls optional sign or phase post-processing.
167    ///
168    /// # Examples
169    ///
170    /// ```rust
171    /// use tenferro_cpu::CpuBackend;
172    /// use tenferro_linalg::{LinalgBackend, QrGauge, QrOptions};
173    /// use tenferro_tensor::Tensor;
174    ///
175    /// let input = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0])?;
176    /// let mut backend = CpuBackend::new();
177    /// let outputs = backend.qr_with_options(
178    ///     &input,
179    ///     QrOptions::default().gauge(QrGauge::PositiveDiagonal),
180    /// )?;
181    /// assert_eq!(outputs[0].shape(), &[2, 2]);
182    /// # Ok::<(), tenferro_tensor::Error>(())
183    /// ```
184    fn qr_with_options(
185        &mut self,
186        input: &Tensor,
187        options: QrOptions,
188    ) -> tenferro_tensor::Result<Vec<Tensor>> {
189        let mut outputs = self.qr(input)?;
190        apply_qr_gauge(options.gauge, &mut outputs)?;
191        Ok(outputs)
192    }
193
194    /// Compute public QR outputs `(Q, R)` from a borrowed tensor view.
195    ///
196    /// Backends may canonicalize the view inside the same placement family, but
197    /// must not silently transfer between CPU and GPU memory.
198    ///
199    /// # Examples
200    ///
201    /// ```rust
202    /// use tenferro_linalg::LinalgBackend;
203    /// use tenferro_cpu::CpuBackend;
204    /// use tenferro_tensor::{TensorView, TypedTensor};
205    ///
206    /// let input = TypedTensor::<f64>::from_vec_col_major(
207    ///     vec![2, 2],
208    ///     vec![1.0, 0.0, 0.0, 2.0],
209    /// )?;
210    /// let mut backend = CpuBackend::new();
211    /// let outputs = backend.qr_read(TensorView::F64(input.as_view()))?;
212    /// assert_eq!(outputs[0].shape(), &[2, 2]);
213    /// assert_eq!(outputs[1].shape(), &[2, 2]);
214    /// # Ok::<(), tenferro_tensor::Error>(())
215    /// ```
216    fn qr_read(&mut self, _input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
217        Err(tenferro_tensor::Error::backend_failure(
218            "qr",
219            "backend does not accept borrowed tensor views at this execution boundary",
220        ))
221    }
222
223    /// Compute public Hermitian eigendecomposition outputs `(values, vectors)`.
224    ///
225    /// The returned vector order is `[values, vectors]`, where `values` has
226    /// shape `[n]` and `vectors` has shape `[n, n]`.
227    fn eigh(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>>;
228
229    /// Compute public Hermitian eigendecomposition outputs with explicit options.
230    ///
231    /// `derivative_eps` is validated for API consistency, but concrete backend
232    /// execution does not perform AD. `gauge` controls optional eigenvector
233    /// post-processing.
234    ///
235    /// # Examples
236    ///
237    /// ```rust
238    /// use tenferro_cpu::CpuBackend;
239    /// use tenferro_linalg::{EighGauge, EighOptions, LinalgBackend};
240    /// use tenferro_tensor::Tensor;
241    ///
242    /// let input = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 2.0])?;
243    /// let mut backend = CpuBackend::new();
244    /// let outputs = backend.eigh_with_options(
245    ///     &input,
246    ///     EighOptions::default()
247    ///         .gauge(EighGauge::CanonicalPivot)
248    ///         .derivative_eps(1.0e-10),
249    /// )?;
250    /// assert_eq!(outputs[0].shape(), &[2]);
251    /// # Ok::<(), tenferro_tensor::Error>(())
252    /// ```
253    fn eigh_with_options(
254        &mut self,
255        input: &Tensor,
256        options: EighOptions,
257    ) -> tenferro_tensor::Result<Vec<Tensor>> {
258        validate_derivative_eps("eigh_with_options", options.derivative_eps)?;
259        let mut outputs = self.eigh(input)?;
260        apply_eigh_gauge(options.gauge, &mut outputs)?;
261        Ok(outputs)
262    }
263
264    /// Compute public Hermitian eigendecomposition outputs from a borrowed tensor view.
265    ///
266    /// Backends may canonicalize the view inside the same placement family, but
267    /// must not silently transfer between CPU and GPU memory.
268    ///
269    /// # Examples
270    ///
271    /// ```rust
272    /// use tenferro_linalg::LinalgBackend;
273    /// use tenferro_cpu::CpuBackend;
274    /// use tenferro_tensor::{TensorView, TypedTensor};
275    ///
276    /// let input = TypedTensor::<f64>::from_vec_col_major(
277    ///     vec![2, 2],
278    ///     vec![1.0, 0.0, 0.0, 2.0],
279    /// )?;
280    /// let mut backend = CpuBackend::new();
281    /// let outputs = backend.eigh_read(TensorView::F64(input.as_view()))?;
282    /// assert_eq!(outputs[0].shape(), &[2]);
283    /// assert_eq!(outputs[1].shape(), &[2, 2]);
284    /// # Ok::<(), tenferro_tensor::Error>(())
285    /// ```
286    fn eigh_read(&mut self, _input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
287        Err(tenferro_tensor::Error::backend_failure(
288            "eigh",
289            "backend does not accept borrowed tensor views at this execution boundary",
290        ))
291    }
292
293    /// Compute Cholesky factorization from a borrowed tensor view.
294    ///
295    /// Backends may canonicalize the view inside the same placement family, but
296    /// must not silently transfer between CPU and GPU memory.
297    ///
298    /// # Examples
299    ///
300    /// ```rust
301    /// use tenferro_linalg::LinalgBackend;
302    /// use tenferro_cpu::CpuBackend;
303    /// use tenferro_tensor::{TensorView, TypedTensor};
304    ///
305    /// let input = TypedTensor::<f64>::from_vec_col_major(
306    ///     vec![2, 2],
307    ///     vec![4.0, 2.0, 2.0, 3.0],
308    /// )?;
309    /// let mut backend = CpuBackend::new();
310    /// let output = backend.cholesky_read(TensorView::F64(input.as_view()))?;
311    /// assert_eq!(output.shape(), &[2, 2]);
312    /// # Ok::<(), tenferro_tensor::Error>(())
313    /// ```
314    fn cholesky_read(&mut self, _input: TensorView<'_>) -> tenferro_tensor::Result<Tensor> {
315        Err(tenferro_tensor::Error::backend_failure(
316            "cholesky",
317            "backend does not accept borrowed tensor views at this execution boundary",
318        ))
319    }
320
321    /// Compute public LU outputs from a borrowed tensor view.
322    ///
323    /// Backends may canonicalize the view inside the same placement family, but
324    /// must not silently transfer between CPU and GPU memory.
325    ///
326    /// # Examples
327    ///
328    /// ```rust
329    /// use tenferro_linalg::LinalgBackend;
330    /// use tenferro_cpu::CpuBackend;
331    /// use tenferro_tensor::{TensorView, TypedTensor};
332    ///
333    /// let input = TypedTensor::<f64>::from_vec_col_major(
334    ///     vec![2, 2],
335    ///     vec![1.0, 3.0, 2.0, 4.0],
336    /// )?;
337    /// let mut backend = CpuBackend::new();
338    /// let outputs = backend.lu_read(TensorView::F64(input.as_view()))?;
339    /// assert_eq!(outputs.len(), 4);
340    /// # Ok::<(), tenferro_tensor::Error>(())
341    /// ```
342    fn lu_read(&mut self, _input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
343        Err(tenferro_tensor::Error::backend_failure(
344            "lu",
345            "backend does not accept borrowed tensor views at this execution boundary",
346        ))
347    }
348
349    /// Compute public full-pivoting LU outputs from a borrowed tensor view.
350    ///
351    /// Backends may canonicalize the view inside the same placement family, but
352    /// must not silently transfer between CPU and GPU memory.
353    ///
354    /// # Examples
355    ///
356    /// ```rust
357    /// use tenferro_linalg::LinalgBackend;
358    /// use tenferro_cpu::CpuBackend;
359    /// use tenferro_tensor::{TensorView, TypedTensor};
360    ///
361    /// let input = TypedTensor::<f64>::from_vec_col_major(
362    ///     vec![2, 2],
363    ///     vec![1.0, 3.0, 2.0, 4.0],
364    /// )?;
365    /// let mut backend = CpuBackend::new();
366    /// let outputs = backend.full_piv_lu_read(TensorView::F64(input.as_view()))?;
367    /// assert_eq!(outputs.len(), 5);
368    /// # Ok::<(), tenferro_tensor::Error>(())
369    /// ```
370    fn full_piv_lu_read(&mut self, _input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
371        Err(tenferro_tensor::Error::backend_failure(
372            "full_piv_lu",
373            "backend does not accept borrowed tensor views at this execution boundary",
374        ))
375    }
376
377    /// Compute general eigendecomposition outputs from a borrowed tensor view.
378    ///
379    /// Backends may canonicalize the view inside the same placement family, but
380    /// must not silently transfer between CPU and GPU memory.
381    ///
382    /// # Examples
383    ///
384    /// ```rust
385    /// use tenferro_linalg::LinalgBackend;
386    /// use tenferro_cpu::CpuBackend;
387    /// use tenferro_tensor::{TensorView, TypedTensor};
388    ///
389    /// let input = TypedTensor::<f64>::from_vec_col_major(
390    ///     vec![2, 2],
391    ///     vec![2.0, 0.0, 0.0, 3.0],
392    /// )?;
393    /// let mut backend = CpuBackend::new();
394    /// let outputs = backend.eig_read(TensorView::F64(input.as_view()))?;
395    /// assert_eq!(outputs.len(), 2);
396    /// # Ok::<(), tenferro_tensor::Error>(())
397    /// ```
398    fn eig_read(&mut self, _input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
399        Err(tenferro_tensor::Error::backend_failure(
400            "eig",
401            "backend does not accept borrowed tensor views at this execution boundary",
402        ))
403    }
404
405    #[doc(hidden)]
406    fn eigh_values(&mut self, _input: &Tensor) -> tenferro_tensor::Result<Tensor> {
407        Err(tenferro_tensor::Error::backend_failure(
408            "eigh_values",
409            format!(
410                "backend {} does not implement internal Hermitian eigenvalues-only decomposition",
411                std::any::type_name::<Self>()
412            ),
413        ))
414    }
415
416    /// Compute public general eigendecomposition outputs `(values, vectors)`.
417    fn eig(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>>;
418
419    #[doc(hidden)]
420    fn eig_values(&mut self, _input: &Tensor) -> tenferro_tensor::Result<Tensor> {
421        Err(tenferro_tensor::Error::backend_failure(
422            "eig_values",
423            format!(
424                "backend {} does not implement internal general eigenvalues-only decomposition",
425                std::any::type_name::<Self>()
426            ),
427        ))
428    }
429
430    /// Solve a dense linear system.
431    fn solve(&mut self, a: &Tensor, b: &Tensor) -> tenferro_tensor::Result<Tensor>;
432
433    #[doc(hidden)]
434    fn lu_solve_prepared(
435        &mut self,
436        _a: &Tensor,
437        _packed_lu: &Tensor,
438        _pivots: &Tensor,
439        _b: &Tensor,
440        _transpose_a: bool,
441        _conjugate_a: bool,
442    ) -> tenferro_tensor::Result<Tensor> {
443        Err(tenferro_tensor::Error::backend_failure(
444            "lu_solve_prepared",
445            format!(
446                "backend {} does not implement internal prepared LU solve",
447                std::any::type_name::<Self>()
448            ),
449        ))
450    }
451}