1use crate::backend::{unsupported_dtype, LinalgBackend};
2
3use super::linalg;
4
5use num_complex::{Complex32, Complex64};
6use tenferro_cpu::{CpuBackend, CpuBackendKind};
7use tenferro_tensor::{
8 validate::validate_nonsingular_u, DType, Error, Tensor, TensorElementwise, TensorStructural,
9 TensorView, TensorViewCanonicalization, TypedTensor,
10};
11
12impl LinalgBackend for CpuBackend {
13 fn cholesky(&mut self, input: &Tensor) -> tenferro_tensor::Result<Tensor> {
14 ensure_host_tensor("cholesky", input)?;
15 match self.kind() {
16 CpuBackendKind::Faer => {
17 #[cfg(feature = "cpu-faer")]
18 {
19 let ctx = self.linalg_context();
20 self.with_linalg_pool(|buffers| match input {
21 Tensor::F32(t) => {
22 linalg::faer::cholesky(ctx.as_ref(), buffers, t).map(Tensor::F32)
23 }
24 Tensor::F64(t) => {
25 linalg::faer::cholesky(ctx.as_ref(), buffers, t).map(Tensor::F64)
26 }
27 Tensor::C32(t) => {
28 linalg::faer::cholesky(ctx.as_ref(), buffers, t).map(Tensor::C32)
29 }
30 Tensor::C64(t) => {
31 linalg::faer::cholesky(ctx.as_ref(), buffers, t).map(Tensor::C64)
32 }
33 _ => Err(unsupported_dtype("cholesky", input.dtype())),
34 })
35 }
36 #[cfg(not(feature = "cpu-faer"))]
37 {
38 Err(unsupported_provider("cholesky", self.kind()))
39 }
40 }
41 CpuBackendKind::Blas => {
42 #[cfg(feature = "cpu-blas")]
43 {
44 self.with_linalg_pool(|buffers| match input {
45 Tensor::F32(t) => linalg::blas::cholesky(buffers, t).map(Tensor::F32),
46 Tensor::F64(t) => linalg::blas::cholesky(buffers, t).map(Tensor::F64),
47 Tensor::C32(t) => linalg::blas::cholesky(buffers, t).map(Tensor::C32),
48 Tensor::C64(t) => linalg::blas::cholesky(buffers, t).map(Tensor::C64),
49 _ => Err(unsupported_dtype("cholesky", input.dtype())),
50 })
51 }
52 #[cfg(not(feature = "cpu-blas"))]
53 {
54 Err(unsupported_provider("cholesky", self.kind()))
55 }
56 }
57 }
58 }
59
60 fn triangular_solve(
61 &mut self,
62 a: &Tensor,
63 b: &Tensor,
64 left_side: bool,
65 lower: bool,
66 transpose_a: bool,
67 unit_diagonal: bool,
68 ) -> tenferro_tensor::Result<Tensor> {
69 ensure_host_tensor("triangular_solve", a)?;
70 ensure_host_tensor("triangular_solve", b)?;
71 match self.kind() {
72 CpuBackendKind::Faer => {
73 #[cfg(feature = "cpu-faer")]
74 {
75 let ctx = self.linalg_context();
76 self.with_linalg_pool(|buffers| match (a, b) {
77 (Tensor::F32(a), Tensor::F32(b)) => linalg::faer::triangular_solve(
78 ctx.as_ref(),
79 buffers,
80 a,
81 b,
82 left_side,
83 lower,
84 transpose_a,
85 unit_diagonal,
86 )
87 .map(Tensor::F32),
88 (Tensor::F64(a), Tensor::F64(b)) => linalg::faer::triangular_solve(
89 ctx.as_ref(),
90 buffers,
91 a,
92 b,
93 left_side,
94 lower,
95 transpose_a,
96 unit_diagonal,
97 )
98 .map(Tensor::F64),
99 (Tensor::C32(a), Tensor::C32(b)) => linalg::faer::triangular_solve(
100 ctx.as_ref(),
101 buffers,
102 a,
103 b,
104 left_side,
105 lower,
106 transpose_a,
107 unit_diagonal,
108 )
109 .map(Tensor::C32),
110 (Tensor::C64(a), Tensor::C64(b)) => linalg::faer::triangular_solve(
111 ctx.as_ref(),
112 buffers,
113 a,
114 b,
115 left_side,
116 lower,
117 transpose_a,
118 unit_diagonal,
119 )
120 .map(Tensor::C64),
121 _ => unsupported_pair("triangular_solve", a, b),
122 })
123 }
124 #[cfg(not(feature = "cpu-faer"))]
125 {
126 Err(unsupported_provider("triangular_solve", self.kind()))
127 }
128 }
129 CpuBackendKind::Blas => {
130 #[cfg(feature = "cpu-blas")]
131 {
132 self.with_linalg_pool(|buffers| match (a, b) {
133 (Tensor::F32(a), Tensor::F32(b)) => linalg::blas::triangular_solve(
134 buffers,
135 a,
136 b,
137 left_side,
138 lower,
139 transpose_a,
140 unit_diagonal,
141 )
142 .map(Tensor::F32),
143 (Tensor::F64(a), Tensor::F64(b)) => linalg::blas::triangular_solve(
144 buffers,
145 a,
146 b,
147 left_side,
148 lower,
149 transpose_a,
150 unit_diagonal,
151 )
152 .map(Tensor::F64),
153 (Tensor::C32(a), Tensor::C32(b)) => linalg::blas::triangular_solve(
154 buffers,
155 a,
156 b,
157 left_side,
158 lower,
159 transpose_a,
160 unit_diagonal,
161 )
162 .map(Tensor::C32),
163 (Tensor::C64(a), Tensor::C64(b)) => linalg::blas::triangular_solve(
164 buffers,
165 a,
166 b,
167 left_side,
168 lower,
169 transpose_a,
170 unit_diagonal,
171 )
172 .map(Tensor::C64),
173 _ => unsupported_pair("triangular_solve", a, b),
174 })
175 }
176 #[cfg(not(feature = "cpu-blas"))]
177 {
178 Err(unsupported_provider("triangular_solve", self.kind()))
179 }
180 }
181 }
182 }
183
184 fn lu(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>> {
185 ensure_host_tensor("lu", input)?;
186 match self.kind() {
187 CpuBackendKind::Faer => {
188 #[cfg(feature = "cpu-faer")]
189 {
190 let ctx = self.linalg_context();
191 self.with_linalg_pool(|buffers| match input {
192 Tensor::F32(t) => linalg::faer::lu(ctx.as_ref(), buffers, t)
193 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
194 Tensor::F64(t) => linalg::faer::lu(ctx.as_ref(), buffers, t)
195 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
196 Tensor::C32(t) => linalg::faer::lu(ctx.as_ref(), buffers, t)
197 .map(|outputs| outputs.into_iter().map(Tensor::C32).collect()),
198 Tensor::C64(t) => linalg::faer::lu(ctx.as_ref(), buffers, t)
199 .map(|outputs| outputs.into_iter().map(Tensor::C64).collect()),
200 _ => Err(unsupported_dtype("lu", input.dtype())),
201 })
202 }
203 #[cfg(not(feature = "cpu-faer"))]
204 {
205 Err(unsupported_provider("lu", self.kind()))
206 }
207 }
208 CpuBackendKind::Blas => {
209 #[cfg(feature = "cpu-blas")]
210 {
211 self.with_linalg_pool(|buffers| match input {
212 Tensor::F32(t) => linalg::blas::lu(buffers, t)
213 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
214 Tensor::F64(t) => linalg::blas::lu(buffers, t)
215 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
216 Tensor::C32(t) => linalg::blas::lu(buffers, t)
217 .map(|outputs| outputs.into_iter().map(Tensor::C32).collect()),
218 Tensor::C64(t) => linalg::blas::lu(buffers, t)
219 .map(|outputs| outputs.into_iter().map(Tensor::C64).collect()),
220 _ => Err(unsupported_dtype("lu", input.dtype())),
221 })
222 }
223 #[cfg(not(feature = "cpu-blas"))]
224 {
225 Err(unsupported_provider("lu", self.kind()))
226 }
227 }
228 }
229 }
230
231 fn lu_factor(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>> {
232 ensure_host_tensor("lu_factor", input)?;
233 match self.kind() {
234 CpuBackendKind::Faer => {
235 #[cfg(feature = "cpu-faer")]
236 {
237 let ctx = self.linalg_context();
238 self.with_linalg_pool(|buffers| match input {
239 Tensor::F32(t) => linalg::faer::lu_factor(ctx.as_ref(), buffers, t).map(
240 |(lu, pivots, parity)| {
241 vec![Tensor::F32(lu), Tensor::I32(pivots), Tensor::F32(parity)]
242 },
243 ),
244 Tensor::F64(t) => linalg::faer::lu_factor(ctx.as_ref(), buffers, t).map(
245 |(lu, pivots, parity)| {
246 vec![Tensor::F64(lu), Tensor::I32(pivots), Tensor::F64(parity)]
247 },
248 ),
249 Tensor::C32(t) => linalg::faer::lu_factor(ctx.as_ref(), buffers, t).map(
250 |(lu, pivots, parity)| {
251 vec![Tensor::C32(lu), Tensor::I32(pivots), Tensor::C32(parity)]
252 },
253 ),
254 Tensor::C64(t) => linalg::faer::lu_factor(ctx.as_ref(), buffers, t).map(
255 |(lu, pivots, parity)| {
256 vec![Tensor::C64(lu), Tensor::I32(pivots), Tensor::C64(parity)]
257 },
258 ),
259 _ => Err(unsupported_dtype("lu_factor", input.dtype())),
260 })
261 }
262 #[cfg(not(feature = "cpu-faer"))]
263 {
264 Err(unsupported_provider("lu_factor", self.kind()))
265 }
266 }
267 CpuBackendKind::Blas => {
268 #[cfg(feature = "cpu-blas")]
269 {
270 self.with_linalg_pool(|buffers| match input {
271 Tensor::F32(t) => {
272 linalg::blas::lu_factor(buffers, t).map(|(lu, pivots, parity)| {
273 vec![Tensor::F32(lu), Tensor::I32(pivots), Tensor::F32(parity)]
274 })
275 }
276 Tensor::F64(t) => {
277 linalg::blas::lu_factor(buffers, t).map(|(lu, pivots, parity)| {
278 vec![Tensor::F64(lu), Tensor::I32(pivots), Tensor::F64(parity)]
279 })
280 }
281 Tensor::C32(t) => {
282 linalg::blas::lu_factor(buffers, t).map(|(lu, pivots, parity)| {
283 vec![Tensor::C32(lu), Tensor::I32(pivots), Tensor::C32(parity)]
284 })
285 }
286 Tensor::C64(t) => {
287 linalg::blas::lu_factor(buffers, t).map(|(lu, pivots, parity)| {
288 vec![Tensor::C64(lu), Tensor::I32(pivots), Tensor::C64(parity)]
289 })
290 }
291 _ => Err(unsupported_dtype("lu_factor", input.dtype())),
292 })
293 }
294 #[cfg(not(feature = "cpu-blas"))]
295 {
296 Err(unsupported_provider("lu_factor", self.kind()))
297 }
298 }
299 }
300 }
301
302 fn full_piv_lu(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>> {
303 ensure_host_tensor("full_piv_lu", input)?;
304 match self.kind() {
305 CpuBackendKind::Faer => {
306 #[cfg(feature = "cpu-faer")]
307 {
308 let ctx = self.linalg_context();
309 self.with_linalg_pool(|buffers| match input {
310 Tensor::F32(t) => linalg::faer::full_piv_lu(ctx.as_ref(), buffers, t)
311 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
312 Tensor::F64(t) => linalg::faer::full_piv_lu(ctx.as_ref(), buffers, t)
313 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
314 Tensor::C32(t) => linalg::faer::full_piv_lu(ctx.as_ref(), buffers, t)
315 .and_then(full_piv_lu_c32_outputs_to_public_tensors),
316 Tensor::C64(t) => linalg::faer::full_piv_lu(ctx.as_ref(), buffers, t)
317 .and_then(full_piv_lu_c64_outputs_to_public_tensors),
318 _ => Err(unsupported_dtype("full_piv_lu", input.dtype())),
319 })
320 }
321 #[cfg(not(feature = "cpu-faer"))]
322 {
323 Err(unsupported_provider("full_piv_lu", self.kind()))
324 }
325 }
326 CpuBackendKind::Blas => {
327 #[cfg(feature = "cpu-blas")]
328 {
329 self.with_linalg_pool(|buffers| match input {
330 Tensor::F32(t) => linalg::blas::full_piv_lu(buffers, t)
331 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
332 Tensor::F64(t) => linalg::blas::full_piv_lu(buffers, t)
333 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
334 Tensor::C32(t) => linalg::blas::full_piv_lu(buffers, t)
335 .and_then(full_piv_lu_c32_outputs_to_public_tensors),
336 Tensor::C64(t) => linalg::blas::full_piv_lu(buffers, t)
337 .and_then(full_piv_lu_c64_outputs_to_public_tensors),
338 _ => Err(unsupported_dtype("full_piv_lu", input.dtype())),
339 })
340 }
341 #[cfg(not(feature = "cpu-blas"))]
342 {
343 Err(unsupported_provider("full_piv_lu", self.kind()))
344 }
345 }
346 }
347 }
348
349 fn full_piv_lu_solve(
350 &mut self,
351 a: &Tensor,
352 b: &Tensor,
353 transpose_a: bool,
354 ) -> tenferro_tensor::Result<Tensor> {
355 ensure_host_tensor("full_piv_lu_solve", a)?;
356 ensure_host_tensor("full_piv_lu_solve", b)?;
357 ensure_supported_linalg_pair("full_piv_lu_solve", a, b)?;
358 if has_zero_dim(a.shape()) || has_zero_dim(b.shape()) {
359 return zeros_like_tensor(b);
360 }
361
362 let (rhs, restore_shape) = if let Some(matrix_rhs_shape) = batched_vector_rhs_shape(a, b) {
363 (
364 self.reshape(b, &matrix_rhs_shape)?,
365 Some(b.shape().to_vec()),
366 )
367 } else {
368 (b.clone(), None)
369 };
370
371 let result = match self.kind() {
372 CpuBackendKind::Faer => {
373 #[cfg(feature = "cpu-faer")]
374 {
375 let ctx = self.linalg_context();
376 self.with_linalg_pool(|buffers| match (a, &rhs) {
377 (Tensor::F32(a), Tensor::F32(b)) => linalg::faer::full_piv_lu_solve(
378 ctx.as_ref(),
379 buffers,
380 a,
381 b,
382 transpose_a,
383 )
384 .map(Tensor::F32),
385 (Tensor::F64(a), Tensor::F64(b)) => linalg::faer::full_piv_lu_solve(
386 ctx.as_ref(),
387 buffers,
388 a,
389 b,
390 transpose_a,
391 )
392 .map(Tensor::F64),
393 (Tensor::C32(a), Tensor::C32(b)) => linalg::faer::full_piv_lu_solve(
394 ctx.as_ref(),
395 buffers,
396 a,
397 b,
398 transpose_a,
399 )
400 .map(Tensor::C32),
401 (Tensor::C64(a), Tensor::C64(b)) => linalg::faer::full_piv_lu_solve(
402 ctx.as_ref(),
403 buffers,
404 a,
405 b,
406 transpose_a,
407 )
408 .map(Tensor::C64),
409 _ => unsupported_pair("full_piv_lu_solve", a, &rhs),
410 })
411 }
412 #[cfg(not(feature = "cpu-faer"))]
413 {
414 Err(unsupported_provider("full_piv_lu_solve", self.kind()))
415 }
416 }
417 CpuBackendKind::Blas => {
418 #[cfg(feature = "cpu-blas")]
419 {
420 self.with_linalg_pool(|buffers| match (a, &rhs) {
421 (Tensor::F32(a), Tensor::F32(b)) => {
422 linalg::blas::full_piv_lu_solve(buffers, a, b, transpose_a)
423 .map(Tensor::F32)
424 }
425 (Tensor::F64(a), Tensor::F64(b)) => {
426 linalg::blas::full_piv_lu_solve(buffers, a, b, transpose_a)
427 .map(Tensor::F64)
428 }
429 (Tensor::C32(a), Tensor::C32(b)) => {
430 linalg::blas::full_piv_lu_solve(buffers, a, b, transpose_a)
431 .map(Tensor::C32)
432 }
433 (Tensor::C64(a), Tensor::C64(b)) => {
434 linalg::blas::full_piv_lu_solve(buffers, a, b, transpose_a)
435 .map(Tensor::C64)
436 }
437 _ => unsupported_pair("full_piv_lu_solve", a, &rhs),
438 })
439 }
440 #[cfg(not(feature = "cpu-blas"))]
441 {
442 Err(unsupported_provider("full_piv_lu_solve", self.kind()))
443 }
444 }
445 }?;
446
447 if let Some(shape) = restore_shape {
448 self.reshape(&result, &shape)
449 } else {
450 Ok(result)
451 }
452 }
453
454 fn svd(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>> {
455 ensure_host_tensor("svd", input)?;
456 match self.kind() {
457 CpuBackendKind::Faer => {
458 #[cfg(feature = "cpu-faer")]
459 {
460 let ctx = self.linalg_context();
461 self.with_linalg_pool(|buffers| match input {
462 Tensor::F32(t) => linalg::faer::svd(ctx.as_ref(), buffers, t)
463 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
464 Tensor::F64(t) => linalg::faer::svd(ctx.as_ref(), buffers, t)
465 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
466 Tensor::C32(t) => linalg::faer::svd(ctx.as_ref(), buffers, t)
467 .and_then(svd_c32_outputs_to_public_tensors),
468 Tensor::C64(t) => linalg::faer::svd(ctx.as_ref(), buffers, t)
469 .and_then(svd_c64_outputs_to_public_tensors),
470 _ => Err(unsupported_dtype("svd", input.dtype())),
471 })
472 }
473 #[cfg(not(feature = "cpu-faer"))]
474 {
475 Err(unsupported_provider("svd", self.kind()))
476 }
477 }
478 CpuBackendKind::Blas => {
479 #[cfg(feature = "cpu-blas")]
480 {
481 self.with_linalg_pool(|buffers| match input {
482 Tensor::F32(t) => linalg::blas::svd(buffers, t)
483 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
484 Tensor::F64(t) => linalg::blas::svd(buffers, t)
485 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
486 Tensor::C32(t) => linalg::blas::svd(buffers, t)
487 .and_then(svd_c32_outputs_to_public_tensors),
488 Tensor::C64(t) => linalg::blas::svd(buffers, t)
489 .and_then(svd_c64_outputs_to_public_tensors),
490 _ => Err(unsupported_dtype("svd", input.dtype())),
491 })
492 }
493 #[cfg(not(feature = "cpu-blas"))]
494 {
495 Err(unsupported_provider("svd", self.kind()))
496 }
497 }
498 }
499 }
500
501 fn svd_values(&mut self, input: &Tensor) -> tenferro_tensor::Result<Tensor> {
502 ensure_host_tensor("svd_values", input)?;
503 match self.kind() {
504 CpuBackendKind::Faer => {
505 #[cfg(feature = "cpu-faer")]
506 {
507 let ctx = self.linalg_context();
508 self.with_linalg_pool(|buffers| match input {
509 Tensor::F32(t) => {
510 linalg::faer::svd_values(ctx.as_ref(), buffers, t).map(Tensor::F32)
511 }
512 Tensor::F64(t) => {
513 linalg::faer::svd_values(ctx.as_ref(), buffers, t).map(Tensor::F64)
514 }
515 Tensor::C32(t) => {
516 linalg::faer::svd_values(ctx.as_ref(), buffers, t).map(Tensor::F32)
517 }
518 Tensor::C64(t) => {
519 linalg::faer::svd_values(ctx.as_ref(), buffers, t).map(Tensor::F64)
520 }
521 _ => Err(unsupported_dtype("svd_values", input.dtype())),
522 })
523 }
524 #[cfg(not(feature = "cpu-faer"))]
525 {
526 Err(unsupported_provider("svd_values", self.kind()))
527 }
528 }
529 CpuBackendKind::Blas => {
530 #[cfg(feature = "cpu-blas")]
531 {
532 self.with_linalg_pool(|buffers| match input {
533 Tensor::F32(t) => linalg::blas::svd_values(buffers, t).map(Tensor::F32),
534 Tensor::F64(t) => linalg::blas::svd_values(buffers, t).map(Tensor::F64),
535 Tensor::C32(t) => linalg::blas::svd_values(buffers, t).map(Tensor::F32),
536 Tensor::C64(t) => linalg::blas::svd_values(buffers, t).map(Tensor::F64),
537 _ => Err(unsupported_dtype("svd_values", input.dtype())),
538 })
539 }
540 #[cfg(not(feature = "cpu-blas"))]
541 {
542 Err(unsupported_provider("svd_values", self.kind()))
543 }
544 }
545 }
546 }
547
548 fn svd_read(&mut self, input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
549 #[cfg(feature = "cpu-faer")]
550 if matches!(self.kind(), CpuBackendKind::Faer) {
551 let can_skip_materialize = match &input {
554 TensorView::F32(view) => linalg::faer::faer_strided_ok(view),
555 TensorView::F64(view) => linalg::faer::faer_strided_ok(view),
556 TensorView::C32(view) => linalg::faer::faer_strided_ok(view),
557 TensorView::C64(view) => linalg::faer::faer_strided_ok(view),
558 _ => false,
559 };
560 if can_skip_materialize {
561 let ctx = self.linalg_context();
562 return self.with_linalg_pool(|buffers| match input {
563 TensorView::F32(view) => linalg::faer::svd_view(ctx.as_ref(), buffers, view)
564 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
565 TensorView::F64(view) => linalg::faer::svd_view(ctx.as_ref(), buffers, view)
566 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
567 TensorView::C32(view) => linalg::faer::svd_view(ctx.as_ref(), buffers, view)
568 .and_then(svd_c32_outputs_to_public_tensors),
569 TensorView::C64(view) => linalg::faer::svd_view(ctx.as_ref(), buffers, view)
570 .and_then(svd_c64_outputs_to_public_tensors),
571 _ => unreachable!("can_skip_materialize only true for supported dtypes"),
572 });
573 }
574 }
575 match input {
578 TensorView::F32(view) => {
579 let compact = self.to_contiguous(&view)?;
580 let input = Tensor::F32(compact);
581 self.svd(&input)
582 }
583 TensorView::F64(view) => {
584 let compact = self.to_contiguous(&view)?;
585 let input = Tensor::F64(compact);
586 self.svd(&input)
587 }
588 TensorView::C32(view) => {
589 let compact = self.to_contiguous(&view)?;
590 let input = Tensor::C32(compact);
591 self.svd(&input)
592 }
593 TensorView::C64(view) => {
594 let compact = self.to_contiguous(&view)?;
595 let input = Tensor::C64(compact);
596 self.svd(&input)
597 }
598 TensorView::I32(_) | TensorView::I64(_) | TensorView::Bool(_) => {
599 Err(unsupported_dtype("svd", input.dtype()))
600 }
601 }
602 }
603
604 fn qr(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>> {
605 ensure_host_tensor("qr", input)?;
606 match self.kind() {
607 CpuBackendKind::Faer => {
608 #[cfg(feature = "cpu-faer")]
609 {
610 let ctx = self.linalg_context();
611 self.with_linalg_pool(|buffers| match input {
612 Tensor::F32(t) => linalg::faer::qr(ctx.as_ref(), buffers, t)
613 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
614 Tensor::F64(t) => linalg::faer::qr(ctx.as_ref(), buffers, t)
615 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
616 Tensor::C32(t) => linalg::faer::qr(ctx.as_ref(), buffers, t)
617 .map(|outputs| outputs.into_iter().map(Tensor::C32).collect()),
618 Tensor::C64(t) => linalg::faer::qr(ctx.as_ref(), buffers, t)
619 .map(|outputs| outputs.into_iter().map(Tensor::C64).collect()),
620 _ => Err(unsupported_dtype("qr", input.dtype())),
621 })
622 }
623 #[cfg(not(feature = "cpu-faer"))]
624 {
625 Err(unsupported_provider("qr", self.kind()))
626 }
627 }
628 CpuBackendKind::Blas => {
629 #[cfg(feature = "cpu-blas")]
630 {
631 self.with_linalg_pool(|buffers| match input {
632 Tensor::F32(t) => linalg::blas::qr(buffers, t)
633 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
634 Tensor::F64(t) => linalg::blas::qr(buffers, t)
635 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
636 Tensor::C32(t) => linalg::blas::qr(buffers, t)
637 .map(|outputs| outputs.into_iter().map(Tensor::C32).collect()),
638 Tensor::C64(t) => linalg::blas::qr(buffers, t)
639 .map(|outputs| outputs.into_iter().map(Tensor::C64).collect()),
640 _ => Err(unsupported_dtype("qr", input.dtype())),
641 })
642 }
643 #[cfg(not(feature = "cpu-blas"))]
644 {
645 Err(unsupported_provider("qr", self.kind()))
646 }
647 }
648 }
649 }
650
651 fn qr_read(&mut self, input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
652 #[cfg(feature = "cpu-faer")]
653 if matches!(self.kind(), CpuBackendKind::Faer) {
654 let can_skip_materialize = match &input {
656 TensorView::F32(view) => linalg::faer::faer_strided_ok(view),
657 TensorView::F64(view) => linalg::faer::faer_strided_ok(view),
658 TensorView::C32(view) => linalg::faer::faer_strided_ok(view),
659 TensorView::C64(view) => linalg::faer::faer_strided_ok(view),
660 _ => false,
661 };
662 if can_skip_materialize {
663 let ctx = self.linalg_context();
664 return self.with_linalg_pool(|buffers| match input {
665 TensorView::F32(view) => linalg::faer::qr_view(ctx.as_ref(), buffers, view)
666 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
667 TensorView::F64(view) => linalg::faer::qr_view(ctx.as_ref(), buffers, view)
668 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
669 TensorView::C32(view) => linalg::faer::qr_view(ctx.as_ref(), buffers, view)
670 .map(|outputs| outputs.into_iter().map(Tensor::C32).collect()),
671 TensorView::C64(view) => linalg::faer::qr_view(ctx.as_ref(), buffers, view)
672 .map(|outputs| outputs.into_iter().map(Tensor::C64).collect()),
673 _ => unreachable!("can_skip_materialize only true for supported dtypes"),
674 });
675 }
676 }
677 match input {
680 TensorView::F32(view) => {
681 let compact = self.to_contiguous(&view)?;
682 let input = Tensor::F32(compact);
683 self.qr(&input)
684 }
685 TensorView::F64(view) => {
686 let compact = self.to_contiguous(&view)?;
687 let input = Tensor::F64(compact);
688 self.qr(&input)
689 }
690 TensorView::C32(view) => {
691 let compact = self.to_contiguous(&view)?;
692 let input = Tensor::C32(compact);
693 self.qr(&input)
694 }
695 TensorView::C64(view) => {
696 let compact = self.to_contiguous(&view)?;
697 let input = Tensor::C64(compact);
698 self.qr(&input)
699 }
700 TensorView::I32(_) | TensorView::I64(_) | TensorView::Bool(_) => {
701 Err(unsupported_dtype("qr", input.dtype()))
702 }
703 }
704 }
705
706 fn eigh(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>> {
707 ensure_host_tensor("eigh", input)?;
708 match self.kind() {
709 CpuBackendKind::Faer => {
710 #[cfg(feature = "cpu-faer")]
711 {
712 let ctx = self.linalg_context();
713 self.with_linalg_pool(|buffers| match input {
714 Tensor::F32(t) => linalg::faer::eigh(ctx.as_ref(), buffers, t)
715 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
716 Tensor::F64(t) => linalg::faer::eigh(ctx.as_ref(), buffers, t)
717 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
718 Tensor::C32(t) => linalg::faer::eigh(ctx.as_ref(), buffers, t)
719 .and_then(eigh_c32_outputs_to_public_tensors),
720 Tensor::C64(t) => linalg::faer::eigh(ctx.as_ref(), buffers, t)
721 .and_then(eigh_c64_outputs_to_public_tensors),
722 _ => Err(unsupported_dtype("eigh", input.dtype())),
723 })
724 }
725 #[cfg(not(feature = "cpu-faer"))]
726 {
727 Err(unsupported_provider("eigh", self.kind()))
728 }
729 }
730 CpuBackendKind::Blas => {
731 #[cfg(feature = "cpu-blas")]
732 {
733 self.with_linalg_pool(|buffers| match input {
734 Tensor::F32(t) => linalg::blas::eigh(buffers, t)
735 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
736 Tensor::F64(t) => linalg::blas::eigh(buffers, t)
737 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
738 Tensor::C32(t) => linalg::blas::eigh(buffers, t)
739 .and_then(eigh_c32_outputs_to_public_tensors),
740 Tensor::C64(t) => linalg::blas::eigh(buffers, t)
741 .and_then(eigh_c64_outputs_to_public_tensors),
742 _ => Err(unsupported_dtype("eigh", input.dtype())),
743 })
744 }
745 #[cfg(not(feature = "cpu-blas"))]
746 {
747 Err(unsupported_provider("eigh", self.kind()))
748 }
749 }
750 }
751 }
752
753 fn eigh_read(&mut self, input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
754 #[cfg(feature = "cpu-faer")]
755 if matches!(self.kind(), CpuBackendKind::Faer) {
756 let can_skip_materialize = match &input {
760 TensorView::F32(view) => linalg::faer::faer_strided_ok(view),
761 TensorView::F64(view) => linalg::faer::faer_strided_ok(view),
762 TensorView::C32(view) => linalg::faer::faer_strided_ok(view),
763 TensorView::C64(view) => linalg::faer::faer_strided_ok(view),
764 _ => false,
765 };
766 if can_skip_materialize {
767 let ctx = self.linalg_context();
768 return self.with_linalg_pool(|buffers| match input {
769 TensorView::F32(view) => linalg::faer::eigh_view(ctx.as_ref(), buffers, view)
770 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
771 TensorView::F64(view) => linalg::faer::eigh_view(ctx.as_ref(), buffers, view)
772 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
773 TensorView::C32(view) => linalg::faer::eigh_view(ctx.as_ref(), buffers, view)
774 .and_then(eigh_c32_outputs_to_public_tensors),
775 TensorView::C64(view) => linalg::faer::eigh_view(ctx.as_ref(), buffers, view)
776 .and_then(eigh_c64_outputs_to_public_tensors),
777 _ => unreachable!("can_skip_materialize only true for supported dtypes"),
778 });
779 }
780 }
781 match input {
784 TensorView::F32(view) => {
785 let compact = self.to_contiguous(&view)?;
786 let input = Tensor::F32(compact);
787 self.eigh(&input)
788 }
789 TensorView::F64(view) => {
790 let compact = self.to_contiguous(&view)?;
791 let input = Tensor::F64(compact);
792 self.eigh(&input)
793 }
794 TensorView::C32(view) => {
795 let compact = self.to_contiguous(&view)?;
796 let input = Tensor::C32(compact);
797 self.eigh(&input)
798 }
799 TensorView::C64(view) => {
800 let compact = self.to_contiguous(&view)?;
801 let input = Tensor::C64(compact);
802 self.eigh(&input)
803 }
804 TensorView::I32(_) | TensorView::I64(_) | TensorView::Bool(_) => {
805 Err(unsupported_dtype("eigh", input.dtype()))
806 }
807 }
808 }
809
810 fn cholesky_read(&mut self, input: TensorView<'_>) -> tenferro_tensor::Result<Tensor> {
811 #[cfg(feature = "cpu-faer")]
812 if matches!(self.kind(), CpuBackendKind::Faer) {
813 let can_skip_materialize = match &input {
814 TensorView::F32(view) => linalg::faer::faer_strided_ok(view),
815 TensorView::F64(view) => linalg::faer::faer_strided_ok(view),
816 TensorView::C32(view) => linalg::faer::faer_strided_ok(view),
817 TensorView::C64(view) => linalg::faer::faer_strided_ok(view),
818 _ => false,
819 };
820 if can_skip_materialize {
821 let ctx = self.linalg_context();
822 return self.with_linalg_pool(|buffers| match input {
823 TensorView::F32(view) => {
824 linalg::faer::cholesky_view(ctx.as_ref(), buffers, view).map(Tensor::F32)
825 }
826 TensorView::F64(view) => {
827 linalg::faer::cholesky_view(ctx.as_ref(), buffers, view).map(Tensor::F64)
828 }
829 TensorView::C32(view) => {
830 linalg::faer::cholesky_view(ctx.as_ref(), buffers, view).map(Tensor::C32)
831 }
832 TensorView::C64(view) => {
833 linalg::faer::cholesky_view(ctx.as_ref(), buffers, view).map(Tensor::C64)
834 }
835 _ => unreachable!("can_skip_materialize only true for supported dtypes"),
836 });
837 }
838 }
839 match input {
840 TensorView::F32(view) => {
841 let compact = self.to_contiguous(&view)?;
842 let input = Tensor::F32(compact);
843 self.cholesky(&input)
844 }
845 TensorView::F64(view) => {
846 let compact = self.to_contiguous(&view)?;
847 let input = Tensor::F64(compact);
848 self.cholesky(&input)
849 }
850 TensorView::C32(view) => {
851 let compact = self.to_contiguous(&view)?;
852 let input = Tensor::C32(compact);
853 self.cholesky(&input)
854 }
855 TensorView::C64(view) => {
856 let compact = self.to_contiguous(&view)?;
857 let input = Tensor::C64(compact);
858 self.cholesky(&input)
859 }
860 TensorView::I32(_) | TensorView::I64(_) | TensorView::Bool(_) => {
861 Err(unsupported_dtype("cholesky", input.dtype()))
862 }
863 }
864 }
865
866 fn lu_read(&mut self, input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
867 #[cfg(feature = "cpu-faer")]
868 if matches!(self.kind(), CpuBackendKind::Faer) {
869 let can_skip_materialize = match &input {
870 TensorView::F32(view) => linalg::faer::faer_strided_ok(view),
871 TensorView::F64(view) => linalg::faer::faer_strided_ok(view),
872 TensorView::C32(view) => linalg::faer::faer_strided_ok(view),
873 TensorView::C64(view) => linalg::faer::faer_strided_ok(view),
874 _ => false,
875 };
876 if can_skip_materialize {
877 let ctx = self.linalg_context();
878 return self.with_linalg_pool(|buffers| match input {
879 TensorView::F32(view) => linalg::faer::lu_view(ctx.as_ref(), buffers, view)
880 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect()),
881 TensorView::F64(view) => linalg::faer::lu_view(ctx.as_ref(), buffers, view)
882 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect()),
883 TensorView::C32(view) => linalg::faer::lu_view(ctx.as_ref(), buffers, view)
884 .map(|outputs| outputs.into_iter().map(Tensor::C32).collect()),
885 TensorView::C64(view) => linalg::faer::lu_view(ctx.as_ref(), buffers, view)
886 .map(|outputs| outputs.into_iter().map(Tensor::C64).collect()),
887 _ => unreachable!("can_skip_materialize only true for supported dtypes"),
888 });
889 }
890 }
891 match input {
892 TensorView::F32(view) => {
893 let compact = self.to_contiguous(&view)?;
894 let input = Tensor::F32(compact);
895 self.lu(&input)
896 }
897 TensorView::F64(view) => {
898 let compact = self.to_contiguous(&view)?;
899 let input = Tensor::F64(compact);
900 self.lu(&input)
901 }
902 TensorView::C32(view) => {
903 let compact = self.to_contiguous(&view)?;
904 let input = Tensor::C32(compact);
905 self.lu(&input)
906 }
907 TensorView::C64(view) => {
908 let compact = self.to_contiguous(&view)?;
909 let input = Tensor::C64(compact);
910 self.lu(&input)
911 }
912 TensorView::I32(_) | TensorView::I64(_) | TensorView::Bool(_) => {
913 Err(unsupported_dtype("lu", input.dtype()))
914 }
915 }
916 }
917
918 fn full_piv_lu_read(&mut self, input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
919 #[cfg(feature = "cpu-faer")]
920 if matches!(self.kind(), CpuBackendKind::Faer) {
921 let can_skip_materialize = match &input {
922 TensorView::F32(view) => linalg::faer::faer_strided_ok(view),
923 TensorView::F64(view) => linalg::faer::faer_strided_ok(view),
924 TensorView::C32(view) => linalg::faer::faer_strided_ok(view),
925 TensorView::C64(view) => linalg::faer::faer_strided_ok(view),
926 _ => false,
927 };
928 if can_skip_materialize {
929 let ctx = self.linalg_context();
930 return self.with_linalg_pool(|buffers| match input {
931 TensorView::F32(view) => {
932 linalg::faer::full_piv_lu_view(ctx.as_ref(), buffers, view)
933 .map(|outputs| outputs.into_iter().map(Tensor::F32).collect())
934 }
935 TensorView::F64(view) => {
936 linalg::faer::full_piv_lu_view(ctx.as_ref(), buffers, view)
937 .map(|outputs| outputs.into_iter().map(Tensor::F64).collect())
938 }
939 TensorView::C32(view) => {
940 linalg::faer::full_piv_lu_view(ctx.as_ref(), buffers, view)
941 .map(|outputs| outputs.into_iter().map(Tensor::C32).collect())
942 }
943 TensorView::C64(view) => {
944 linalg::faer::full_piv_lu_view(ctx.as_ref(), buffers, view)
945 .map(|outputs| outputs.into_iter().map(Tensor::C64).collect())
946 }
947 _ => unreachable!("can_skip_materialize only true for supported dtypes"),
948 });
949 }
950 }
951 match input {
952 TensorView::F32(view) => {
953 let compact = self.to_contiguous(&view)?;
954 let input = Tensor::F32(compact);
955 self.full_piv_lu(&input)
956 }
957 TensorView::F64(view) => {
958 let compact = self.to_contiguous(&view)?;
959 let input = Tensor::F64(compact);
960 self.full_piv_lu(&input)
961 }
962 TensorView::C32(view) => {
963 let compact = self.to_contiguous(&view)?;
964 let input = Tensor::C32(compact);
965 self.full_piv_lu(&input)
966 }
967 TensorView::C64(view) => {
968 let compact = self.to_contiguous(&view)?;
969 let input = Tensor::C64(compact);
970 self.full_piv_lu(&input)
971 }
972 TensorView::I32(_) | TensorView::I64(_) | TensorView::Bool(_) => {
973 Err(unsupported_dtype("full_piv_lu", input.dtype()))
974 }
975 }
976 }
977
978 fn eig_read(&mut self, input: TensorView<'_>) -> tenferro_tensor::Result<Vec<Tensor>> {
979 match input {
981 TensorView::F32(view) => {
982 let compact = self.to_contiguous(&view)?;
983 let input = Tensor::F32(compact);
984 self.eig(&input)
985 }
986 TensorView::F64(view) => {
987 let compact = self.to_contiguous(&view)?;
988 let input = Tensor::F64(compact);
989 self.eig(&input)
990 }
991 TensorView::C32(view) => {
992 let compact = self.to_contiguous(&view)?;
993 let input = Tensor::C32(compact);
994 self.eig(&input)
995 }
996 TensorView::C64(view) => {
997 let compact = self.to_contiguous(&view)?;
998 let input = Tensor::C64(compact);
999 self.eig(&input)
1000 }
1001 TensorView::I32(_) | TensorView::I64(_) | TensorView::Bool(_) => {
1002 Err(unsupported_dtype("eig", input.dtype()))
1003 }
1004 }
1005 }
1006
1007 fn eigh_values(&mut self, input: &Tensor) -> tenferro_tensor::Result<Tensor> {
1008 ensure_host_tensor("eigh_values", input)?;
1009 match self.kind() {
1010 CpuBackendKind::Faer => {
1011 #[cfg(feature = "cpu-faer")]
1012 {
1013 let ctx = self.linalg_context();
1014 self.with_linalg_pool(|buffers| match input {
1015 Tensor::F32(t) => {
1016 linalg::faer::eigh_values(ctx.as_ref(), buffers, t).map(Tensor::F32)
1017 }
1018 Tensor::F64(t) => {
1019 linalg::faer::eigh_values(ctx.as_ref(), buffers, t).map(Tensor::F64)
1020 }
1021 Tensor::C32(t) => {
1022 linalg::faer::eigh_values(ctx.as_ref(), buffers, t).map(Tensor::F32)
1023 }
1024 Tensor::C64(t) => {
1025 linalg::faer::eigh_values(ctx.as_ref(), buffers, t).map(Tensor::F64)
1026 }
1027 _ => Err(unsupported_dtype("eigh_values", input.dtype())),
1028 })
1029 }
1030 #[cfg(not(feature = "cpu-faer"))]
1031 {
1032 Err(unsupported_provider("eigh_values", self.kind()))
1033 }
1034 }
1035 CpuBackendKind::Blas => {
1036 #[cfg(feature = "cpu-blas")]
1037 {
1038 self.with_linalg_pool(|buffers| match input {
1039 Tensor::F32(t) => linalg::blas::eigh_values(buffers, t).map(Tensor::F32),
1040 Tensor::F64(t) => linalg::blas::eigh_values(buffers, t).map(Tensor::F64),
1041 Tensor::C32(t) => linalg::blas::eigh_values(buffers, t).map(Tensor::F32),
1042 Tensor::C64(t) => linalg::blas::eigh_values(buffers, t).map(Tensor::F64),
1043 _ => Err(unsupported_dtype("eigh_values", input.dtype())),
1044 })
1045 }
1046 #[cfg(not(feature = "cpu-blas"))]
1047 {
1048 Err(unsupported_provider("eigh_values", self.kind()))
1049 }
1050 }
1051 }
1052 }
1053
1054 fn eig(&mut self, input: &Tensor) -> tenferro_tensor::Result<Vec<Tensor>> {
1055 ensure_host_tensor("eig", input)?;
1056 if !matches!(
1057 input,
1058 Tensor::F32(_) | Tensor::F64(_) | Tensor::C32(_) | Tensor::C64(_)
1059 ) {
1060 return Err(unsupported_dtype("eig", input.dtype()));
1061 }
1062 match self.kind() {
1063 CpuBackendKind::Faer => {
1064 #[cfg(feature = "cpu-faer")]
1065 {
1066 let ctx = self.linalg_context();
1067 self.with_linalg_pool(|buffers| linalg::faer::eig(ctx.as_ref(), buffers, input))
1068 }
1069 #[cfg(not(feature = "cpu-faer"))]
1070 {
1071 Err(unsupported_provider("eig", self.kind()))
1072 }
1073 }
1074 CpuBackendKind::Blas => {
1075 #[cfg(feature = "cpu-blas")]
1076 {
1077 self.with_linalg_pool(|buffers| linalg::blas::eig(buffers, input))
1078 }
1079 #[cfg(not(feature = "cpu-blas"))]
1080 {
1081 Err(unsupported_provider("eig", self.kind()))
1082 }
1083 }
1084 }
1085 }
1086
1087 fn eig_values(&mut self, input: &Tensor) -> tenferro_tensor::Result<Tensor> {
1088 ensure_host_tensor("eig_values", input)?;
1089 if !matches!(
1090 input,
1091 Tensor::F32(_) | Tensor::F64(_) | Tensor::C32(_) | Tensor::C64(_)
1092 ) {
1093 return Err(unsupported_dtype("eig_values", input.dtype()));
1094 }
1095 match self.kind() {
1096 CpuBackendKind::Faer => {
1097 #[cfg(feature = "cpu-faer")]
1098 {
1099 let ctx = self.linalg_context();
1100 self.with_linalg_pool(|buffers| {
1101 linalg::faer::eig_values(ctx.as_ref(), buffers, input)
1102 })
1103 }
1104 #[cfg(not(feature = "cpu-faer"))]
1105 {
1106 Err(unsupported_provider("eig_values", self.kind()))
1107 }
1108 }
1109 CpuBackendKind::Blas => {
1110 #[cfg(feature = "cpu-blas")]
1111 {
1112 self.with_linalg_pool(|buffers| linalg::blas::eig_values(buffers, input))
1113 }
1114 #[cfg(not(feature = "cpu-blas"))]
1115 {
1116 Err(unsupported_provider("eig_values", self.kind()))
1117 }
1118 }
1119 }
1120 }
1121
1122 fn lu_solve_prepared(
1123 &mut self,
1124 a: &Tensor,
1125 packed_lu: &Tensor,
1126 pivots: &Tensor,
1127 b: &Tensor,
1128 transpose_a: bool,
1129 conjugate_a: bool,
1130 ) -> tenferro_tensor::Result<Tensor> {
1131 const OP: &str = "lu_solve_prepared";
1132
1133 ensure_host_tensor(OP, a)?;
1134 ensure_host_tensor(OP, packed_lu)?;
1135 ensure_host_tensor(OP, pivots)?;
1136 ensure_host_tensor(OP, b)?;
1137 ensure_supported_linalg_pair(OP, a, b)?;
1138 ensure_supported_linalg_pair(OP, a, packed_lu)?;
1139 if !matches!(pivots, Tensor::I32(_)) {
1140 return Err(Error::DTypeMismatch {
1141 op: OP,
1142 lhs: DType::I32,
1143 rhs: pivots.dtype(),
1144 });
1145 }
1146 if has_zero_dim(a.shape()) || has_zero_dim(b.shape()) {
1147 return zeros_like_tensor(b);
1148 }
1149
1150 let (rhs, restore_shape) = if let Some(matrix_rhs_shape) = batched_vector_rhs_shape(a, b) {
1151 (
1152 self.reshape(b, &matrix_rhs_shape)?,
1153 Some(b.shape().to_vec()),
1154 )
1155 } else {
1156 (b.clone(), None)
1157 };
1158
1159 validate_lu_solve_prepared_shapes(packed_lu.shape(), pivots.shape(), rhs.shape())?;
1160 validate_nonsingular_u(packed_lu)?;
1161 let lu_op = if conjugate_a {
1162 self.conj(packed_lu)?
1163 } else {
1164 packed_lu.clone()
1165 };
1166 let result = if transpose_a {
1167 let z = self.triangular_solve(&lu_op, &rhs, true, false, true, false)?;
1168 let y = self.triangular_solve(&lu_op, &z, true, true, true, true)?;
1169 apply_lu_pivots_cpu(&y, pivots, true)?
1170 } else {
1171 let pb = apply_lu_pivots_cpu(&rhs, pivots, false)?;
1172 let y = self.triangular_solve(&lu_op, &pb, true, true, false, true)?;
1173 self.triangular_solve(&lu_op, &y, true, false, false, false)?
1174 };
1175
1176 if let Some(shape) = restore_shape {
1177 self.reshape(&result, &shape)
1178 } else {
1179 Ok(result)
1180 }
1181 }
1182
1183 fn solve(&mut self, a: &Tensor, b: &Tensor) -> tenferro_tensor::Result<Tensor> {
1184 ensure_host_tensor("solve", a)?;
1185 ensure_host_tensor("solve", b)?;
1186 ensure_supported_linalg_pair("solve", a, b)?;
1187 if has_zero_dim(a.shape()) || has_zero_dim(b.shape()) {
1188 return zeros_like_tensor(b);
1189 }
1190
1191 let (rhs, restore_shape) = if let Some(matrix_rhs_shape) = batched_vector_rhs_shape(a, b) {
1192 (
1193 self.reshape(b, &matrix_rhs_shape)?,
1194 Some(b.shape().to_vec()),
1195 )
1196 } else {
1197 (b.clone(), None)
1198 };
1199
1200 let result = match self.kind() {
1201 CpuBackendKind::Faer => {
1202 #[cfg(feature = "cpu-faer")]
1203 {
1204 let ctx = self.linalg_context();
1205 self.with_linalg_pool(|buffers| match (a, &rhs) {
1206 (Tensor::F32(a), Tensor::F32(b)) => {
1207 linalg::faer::solve(ctx.as_ref(), buffers, a, b, false).map(Tensor::F32)
1208 }
1209 (Tensor::F64(a), Tensor::F64(b)) => {
1210 linalg::faer::solve(ctx.as_ref(), buffers, a, b, false).map(Tensor::F64)
1211 }
1212 (Tensor::C32(a), Tensor::C32(b)) => {
1213 linalg::faer::solve(ctx.as_ref(), buffers, a, b, false).map(Tensor::C32)
1214 }
1215 (Tensor::C64(a), Tensor::C64(b)) => {
1216 linalg::faer::solve(ctx.as_ref(), buffers, a, b, false).map(Tensor::C64)
1217 }
1218 _ => unsupported_pair("solve", a, &rhs),
1219 })
1220 }
1221 #[cfg(not(feature = "cpu-faer"))]
1222 {
1223 Err(unsupported_provider("solve", self.kind()))
1224 }
1225 }
1226 CpuBackendKind::Blas => {
1227 #[cfg(feature = "cpu-blas")]
1228 {
1229 self.with_linalg_pool(|buffers| match (a, &rhs) {
1230 (Tensor::F32(a), Tensor::F32(b)) => {
1231 linalg::blas::solve(buffers, a, b, false).map(Tensor::F32)
1232 }
1233 (Tensor::F64(a), Tensor::F64(b)) => {
1234 linalg::blas::solve(buffers, a, b, false).map(Tensor::F64)
1235 }
1236 (Tensor::C32(a), Tensor::C32(b)) => {
1237 linalg::blas::solve(buffers, a, b, false).map(Tensor::C32)
1238 }
1239 (Tensor::C64(a), Tensor::C64(b)) => {
1240 linalg::blas::solve(buffers, a, b, false).map(Tensor::C64)
1241 }
1242 _ => unsupported_pair("solve", a, &rhs),
1243 })
1244 }
1245 #[cfg(not(feature = "cpu-blas"))]
1246 {
1247 Err(unsupported_provider("solve", self.kind()))
1248 }
1249 }
1250 }?;
1251
1252 if let Some(shape) = restore_shape {
1253 self.reshape(&result, &shape)
1254 } else {
1255 Ok(result)
1256 }
1257 }
1258}
1259
1260fn ensure_host_tensor(op: &'static str, input: &Tensor) -> tenferro_tensor::Result<()> {
1261 match input {
1262 Tensor::F32(t) => ensure_host_typed_tensor(op, t),
1263 Tensor::F64(t) => ensure_host_typed_tensor(op, t),
1264 Tensor::I32(t) => ensure_host_typed_tensor(op, t),
1265 Tensor::I64(t) => ensure_host_typed_tensor(op, t),
1266 Tensor::Bool(t) => ensure_host_typed_tensor(op, t),
1267 Tensor::C32(t) => ensure_host_typed_tensor(op, t),
1268 Tensor::C64(t) => ensure_host_typed_tensor(op, t),
1269 }
1270}
1271
1272fn ensure_host_typed_tensor<T: 'static>(
1273 op: &'static str,
1274 input: &TypedTensor<T>,
1275) -> tenferro_tensor::Result<()> {
1276 if input.as_view().backend_buffer().is_some() {
1277 return Err(Error::backend_failure(
1278 op,
1279 "CPU linalg backend received a backend buffer; download the tensor to host before CPU execution",
1280 ));
1281 }
1282 Ok(())
1283}
1284
1285fn ensure_supported_linalg_pair(
1286 op: &'static str,
1287 lhs: &Tensor,
1288 rhs: &Tensor,
1289) -> tenferro_tensor::Result<()> {
1290 if lhs.dtype() != rhs.dtype() {
1291 return Err(Error::DTypeMismatch {
1292 op,
1293 lhs: lhs.dtype(),
1294 rhs: rhs.dtype(),
1295 });
1296 }
1297 match lhs {
1298 Tensor::F32(_) | Tensor::F64(_) | Tensor::C32(_) | Tensor::C64(_) => Ok(()),
1299 Tensor::I32(_) | Tensor::I64(_) | Tensor::Bool(_) => {
1300 Err(unsupported_dtype(op, lhs.dtype()))
1301 }
1302 }
1303}
1304
1305fn has_zero_dim(shape: &[usize]) -> bool {
1306 shape.contains(&0)
1307}
1308
1309fn checked_product(
1310 op: &'static str,
1311 role: &'static str,
1312 shape: &[usize],
1313) -> tenferro_tensor::Result<usize> {
1314 shape.iter().try_fold(1usize, |acc, &dim| {
1315 acc.checked_mul(dim).ok_or_else(|| Error::InvalidConfig {
1316 op,
1317 message: format!("{role} element count overflow"),
1318 })
1319 })
1320}
1321
1322fn batch_count(op: &'static str, batch_shape: &[usize]) -> tenferro_tensor::Result<usize> {
1323 Ok(checked_product(op, "batch shape", batch_shape)?.max(1))
1324}
1325
1326fn checked_batch_offset(
1327 op: &'static str,
1328 role: &'static str,
1329 batch: usize,
1330 stride: usize,
1331) -> tenferro_tensor::Result<usize> {
1332 batch
1333 .checked_mul(stride)
1334 .ok_or_else(|| Error::InvalidConfig {
1335 op,
1336 message: format!("{role} overflows usize"),
1337 })
1338}
1339
1340fn batched_vector_rhs_shape(a: &Tensor, b: &Tensor) -> Option<Vec<usize>> {
1341 if b.shape().len() == 1 {
1342 return Some(vec![b.shape()[0], 1]);
1343 }
1344
1345 let is_batched_vector_rhs = a.shape().len() == b.shape().len() + 1
1346 && !b.shape().is_empty()
1347 && b.shape()[0] == a.shape()[0]
1348 && b.shape()[1..] == a.shape()[2..];
1349 if !is_batched_vector_rhs {
1350 return None;
1351 }
1352
1353 let mut rhs_shape = vec![b.shape()[0], 1];
1354 rhs_shape.extend_from_slice(&b.shape()[1..]);
1355 Some(rhs_shape)
1356}
1357
1358fn zeros_like_tensor(input: &Tensor) -> tenferro_tensor::Result<Tensor> {
1359 Ok(match input {
1360 Tensor::F32(t) => Tensor::F32(TypedTensor::zeros(t.shape().to_vec())?),
1361 Tensor::F64(t) => Tensor::F64(TypedTensor::zeros(t.shape().to_vec())?),
1362 Tensor::I32(t) => Tensor::I32(TypedTensor::zeros(t.shape().to_vec())?),
1363 Tensor::I64(t) => Tensor::I64(TypedTensor::zeros(t.shape().to_vec())?),
1364 Tensor::Bool(t) => Tensor::Bool(TypedTensor::from_vec_col_major(
1365 t.shape().to_vec(),
1366 vec![false; t.n_elements()],
1367 )?),
1368 Tensor::C32(t) => Tensor::C32(TypedTensor::zeros(t.shape().to_vec())?),
1369 Tensor::C64(t) => Tensor::C64(TypedTensor::zeros(t.shape().to_vec())?),
1370 })
1371}
1372
1373fn complex32_real_part_tensor(
1374 values: TypedTensor<Complex32>,
1375) -> tenferro_tensor::Result<TypedTensor<f32>> {
1376 let mut out = TypedTensor::from_vec_col_major(
1377 values.shape().to_vec(),
1378 values.host_data()?.iter().map(|value| value.re).collect(),
1379 )?;
1380 out.set_placement(values.placement().clone());
1381 Ok(out)
1382}
1383
1384fn complex64_real_part_tensor(
1385 values: TypedTensor<Complex64>,
1386) -> tenferro_tensor::Result<TypedTensor<f64>> {
1387 let mut out = TypedTensor::from_vec_col_major(
1388 values.shape().to_vec(),
1389 values.host_data()?.iter().map(|value| value.re).collect(),
1390 )?;
1391 out.set_placement(values.placement().clone());
1392 Ok(out)
1393}
1394
1395fn svd_output_count_error(count: usize) -> Error {
1396 Error::backend_failure("svd", format!("expected 3 outputs, got {count}"))
1397}
1398
1399fn full_piv_lu_output_count_error(count: usize) -> Error {
1400 Error::backend_failure("full_piv_lu", format!("expected 5 outputs, got {count}"))
1401}
1402
1403fn eigh_output_count_error(count: usize) -> Error {
1404 Error::backend_failure("eigh", format!("expected 2 outputs, got {count}"))
1405}
1406
1407fn full_piv_lu_c32_outputs_to_public_tensors(
1408 outputs: Vec<TypedTensor<Complex32>>,
1409) -> tenferro_tensor::Result<Vec<Tensor>> {
1410 let count = outputs.len();
1411 let mut outputs = outputs.into_iter();
1412 match (
1413 outputs.next(),
1414 outputs.next(),
1415 outputs.next(),
1416 outputs.next(),
1417 outputs.next(),
1418 outputs.next(),
1419 ) {
1420 (Some(p), Some(l), Some(u), Some(q), Some(parity), None) => Ok(vec![
1421 Tensor::C32(p),
1422 Tensor::C32(l),
1423 Tensor::C32(u),
1424 Tensor::C32(q),
1425 Tensor::F32(complex32_real_part_tensor(parity)?),
1426 ]),
1427 _ => Err(full_piv_lu_output_count_error(count)),
1428 }
1429}
1430
1431fn full_piv_lu_c64_outputs_to_public_tensors(
1432 outputs: Vec<TypedTensor<Complex64>>,
1433) -> tenferro_tensor::Result<Vec<Tensor>> {
1434 let count = outputs.len();
1435 let mut outputs = outputs.into_iter();
1436 match (
1437 outputs.next(),
1438 outputs.next(),
1439 outputs.next(),
1440 outputs.next(),
1441 outputs.next(),
1442 outputs.next(),
1443 ) {
1444 (Some(p), Some(l), Some(u), Some(q), Some(parity), None) => Ok(vec![
1445 Tensor::C64(p),
1446 Tensor::C64(l),
1447 Tensor::C64(u),
1448 Tensor::C64(q),
1449 Tensor::F64(complex64_real_part_tensor(parity)?),
1450 ]),
1451 _ => Err(full_piv_lu_output_count_error(count)),
1452 }
1453}
1454
1455fn svd_c32_outputs_to_public_tensors(
1456 outputs: Vec<TypedTensor<Complex32>>,
1457) -> tenferro_tensor::Result<Vec<Tensor>> {
1458 let count = outputs.len();
1459 let mut outputs = outputs.into_iter();
1460 match (
1461 outputs.next(),
1462 outputs.next(),
1463 outputs.next(),
1464 outputs.next(),
1465 ) {
1466 (Some(u), Some(values), Some(vt), None) => Ok(vec![
1467 Tensor::C32(u),
1468 Tensor::F32(complex32_real_part_tensor(values)?),
1469 Tensor::C32(vt),
1470 ]),
1471 _ => Err(svd_output_count_error(count)),
1472 }
1473}
1474
1475fn svd_c64_outputs_to_public_tensors(
1476 outputs: Vec<TypedTensor<Complex64>>,
1477) -> tenferro_tensor::Result<Vec<Tensor>> {
1478 let count = outputs.len();
1479 let mut outputs = outputs.into_iter();
1480 match (
1481 outputs.next(),
1482 outputs.next(),
1483 outputs.next(),
1484 outputs.next(),
1485 ) {
1486 (Some(u), Some(values), Some(vt), None) => Ok(vec![
1487 Tensor::C64(u),
1488 Tensor::F64(complex64_real_part_tensor(values)?),
1489 Tensor::C64(vt),
1490 ]),
1491 _ => Err(svd_output_count_error(count)),
1492 }
1493}
1494
1495fn eigh_c32_outputs_to_public_tensors(
1496 outputs: Vec<TypedTensor<Complex32>>,
1497) -> tenferro_tensor::Result<Vec<Tensor>> {
1498 let count = outputs.len();
1499 let mut outputs = outputs.into_iter();
1500 match (outputs.next(), outputs.next(), outputs.next()) {
1501 (Some(values), Some(vectors), None) => Ok(vec![
1502 Tensor::F32(complex32_real_part_tensor(values)?),
1503 Tensor::C32(vectors),
1504 ]),
1505 _ => Err(eigh_output_count_error(count)),
1506 }
1507}
1508
1509fn eigh_c64_outputs_to_public_tensors(
1510 outputs: Vec<TypedTensor<Complex64>>,
1511) -> tenferro_tensor::Result<Vec<Tensor>> {
1512 let count = outputs.len();
1513 let mut outputs = outputs.into_iter();
1514 match (outputs.next(), outputs.next(), outputs.next()) {
1515 (Some(values), Some(vectors), None) => Ok(vec![
1516 Tensor::F64(complex64_real_part_tensor(values)?),
1517 Tensor::C64(vectors),
1518 ]),
1519 _ => Err(eigh_output_count_error(count)),
1520 }
1521}
1522
1523fn apply_lu_pivots_cpu(
1524 input: &Tensor,
1525 pivots: &Tensor,
1526 inverse: bool,
1527) -> tenferro_tensor::Result<Tensor> {
1528 let Tensor::I32(pivots) = pivots else {
1529 return Err(Error::DTypeMismatch {
1530 op: "lu_solve_prepared",
1531 lhs: DType::I32,
1532 rhs: pivots.dtype(),
1533 });
1534 };
1535 match input {
1536 Tensor::F32(t) => apply_lu_pivots_typed(t, pivots, inverse).map(Tensor::F32),
1537 Tensor::F64(t) => apply_lu_pivots_typed(t, pivots, inverse).map(Tensor::F64),
1538 Tensor::C32(t) => apply_lu_pivots_typed(t, pivots, inverse).map(Tensor::C32),
1539 Tensor::C64(t) => apply_lu_pivots_typed(t, pivots, inverse).map(Tensor::C64),
1540 Tensor::I32(_) | Tensor::I64(_) | Tensor::Bool(_) => {
1541 Err(unsupported_dtype("lu_solve_prepared", input.dtype()))
1542 }
1543 }
1544}
1545
1546fn apply_lu_pivots_typed<T: Clone>(
1547 input: &TypedTensor<T>,
1548 pivots: &TypedTensor<i32>,
1549 inverse: bool,
1550) -> tenferro_tensor::Result<TypedTensor<T>> {
1551 let shape = input.shape();
1552 if shape.len() < 2 {
1553 return Err(Error::RankMismatch {
1554 op: "lu_solve_prepared",
1555 expected: 2,
1556 actual: shape.len(),
1557 });
1558 }
1559 let rows = shape[0];
1560 let cols = shape[1];
1561 let k = pivots.shape()[0];
1562 if k > rows || pivots.shape()[1..] != shape[2..] {
1563 return Err(Error::ShapeMismatch {
1564 op: "lu_solve_prepared",
1565 lhs: pivots.shape().to_vec(),
1566 rhs: shape.to_vec(),
1567 });
1568 }
1569 let batch_total = batch_count("lu_solve_prepared", &shape[2..])?;
1570 let matrix_stride = checked_product("lu_solve_prepared", "matrix shape", &[rows, cols])?;
1571 let pivot_stride = k;
1572 let input_data = input.host_data()?;
1573 let pivot_data = pivots.host_data()?;
1574 let mut data = Vec::with_capacity(input_data.len());
1575
1576 for batch in 0..batch_total {
1577 let mut perm: Vec<usize> = (0..rows).collect();
1578 let pivot_offset = checked_batch_offset(
1579 "lu_solve_prepared",
1580 "pivot batch offset",
1581 batch,
1582 pivot_stride,
1583 )?;
1584 for step in 0..k {
1585 let pivot_one_based = pivot_data[pivot_offset + step];
1586 if pivot_one_based <= 0 {
1587 return Err(Error::backend_failure(
1588 "lu_solve_prepared",
1589 "LU pivot index must be 1-based and positive",
1590 ));
1591 }
1592 let pivot = usize::try_from(pivot_one_based - 1).map_err(|_| {
1593 Error::backend_failure("lu_solve_prepared", "LU pivot index is invalid")
1594 })?;
1595 if pivot >= rows {
1596 return Err(Error::backend_failure(
1597 "lu_solve_prepared",
1598 "LU pivot index is out of bounds",
1599 ));
1600 }
1601 perm.swap(step, pivot);
1602 }
1603 let row_map = if inverse {
1604 let mut inv = vec![0usize; rows];
1605 for (row, &source) in perm.iter().enumerate() {
1606 inv[source] = row;
1607 }
1608 inv
1609 } else {
1610 perm
1611 };
1612 let batch_offset = checked_batch_offset(
1613 "lu_solve_prepared",
1614 "matrix batch offset",
1615 batch,
1616 matrix_stride,
1617 )?;
1618 for col in 0..cols {
1619 for &source_row in &row_map {
1620 data.push(input_data[batch_offset + source_row + col * rows].clone());
1621 }
1622 }
1623 }
1624
1625 TypedTensor::from_vec_col_major(shape.to_vec(), data)
1626}
1627
1628fn validate_lu_solve_prepared_shapes(
1629 lu_shape: &[usize],
1630 pivots_shape: &[usize],
1631 b_shape: &[usize],
1632) -> tenferro_tensor::Result<()> {
1633 let n = square_matrix_dim("lu_solve_prepared", lu_shape)?;
1634 let (b_rows, _) = matrix_dims("lu_solve_prepared", b_shape)?;
1635 if b_rows != n {
1636 return Err(Error::InvalidConfig {
1637 op: "lu_solve_prepared",
1638 message: format!("rhs row count mismatch: expected {n}, got {b_rows}"),
1639 });
1640 }
1641 if lu_shape[2..] != b_shape[2..] {
1642 return Err(Error::ShapeMismatch {
1643 op: "lu_solve_prepared",
1644 lhs: lu_shape.to_vec(),
1645 rhs: b_shape.to_vec(),
1646 });
1647 }
1648 let mut expected_pivots = vec![n];
1649 expected_pivots.extend_from_slice(&lu_shape[2..]);
1650 if pivots_shape != expected_pivots {
1651 return Err(Error::ShapeMismatch {
1652 op: "lu_solve_prepared",
1653 lhs: expected_pivots,
1654 rhs: pivots_shape.to_vec(),
1655 });
1656 }
1657 Ok(())
1658}
1659
1660fn matrix_dims(op: &'static str, shape: &[usize]) -> tenferro_tensor::Result<(usize, usize)> {
1661 if shape.len() < 2 {
1662 return Err(Error::RankMismatch {
1663 op,
1664 expected: 2,
1665 actual: shape.len(),
1666 });
1667 }
1668 Ok((shape[0], shape[1]))
1669}
1670
1671fn square_matrix_dim(op: &'static str, shape: &[usize]) -> tenferro_tensor::Result<usize> {
1672 let (rows, cols) = matrix_dims(op, shape)?;
1673 if rows != cols {
1674 return Err(Error::ShapeMismatch {
1675 op,
1676 lhs: vec![rows],
1677 rhs: vec![cols],
1678 });
1679 }
1680 Ok(rows)
1681}
1682
1683#[allow(dead_code)]
1686fn unsupported_provider(op: &'static str, kind: CpuBackendKind) -> Error {
1687 Error::InvalidConfig {
1688 op,
1689 message: format!("CPU linalg provider {kind:?} is not compiled in"),
1690 }
1691}
1692
1693fn unsupported_pair(
1694 op: &'static str,
1695 lhs: &Tensor,
1696 rhs: &Tensor,
1697) -> tenferro_tensor::Result<Tensor> {
1698 if lhs.dtype() != rhs.dtype() {
1699 Err(Error::DTypeMismatch {
1700 op,
1701 lhs: lhs.dtype(),
1702 rhs: rhs.dtype(),
1703 })
1704 } else {
1705 Err(unsupported_dtype(op, lhs.dtype()))
1706 }
1707}