tenferro_linalg/
extension.rs

1use std::any::Any;
2use std::hash::Hasher;
3use std::sync::Arc;
4
5use num_complex::{Complex32, Complex64};
6use tenferro_cpu::with_cpu_exec_session;
7use tenferro_extension_macros::define_extension_runtime;
8use tenferro_ops::SymDim;
9use tenferro_runtime::extension::{ExtensionExecutionContext, ExtensionOp};
10use tenferro_tensor::{BackendSession, DType, Error, ErrorKind, Tensor, TensorBackend, TensorRead};
11
12#[cfg(feature = "cuda")]
13use tenferro_gpu::cuda::with_cuda_exec_session;
14
15use crate::backend::LinalgBackend;
16
17mod gauge;
18#[cfg(all(test, not(feature = "cuda")))]
19mod tests;
20
21pub(crate) use gauge::{apply_eigh_gauge, apply_qr_gauge};
22
23pub const LINALG_EXTENSION_FAMILY_ID: &str = "tenferro-linalg.linalg.v1";
24
25/// Default derivative regularization used by decomposition AD rules.
26///
27/// This epsilon is used only when differentiating decomposition formulas with
28/// repeated or nearly repeated spectral values. It is not a solver tolerance.
29///
30/// # Examples
31///
32/// ```rust
33/// use tenferro_linalg::{SvdOptions, DEFAULT_DECOMPOSITION_DERIVATIVE_EPS};
34///
35/// let options = SvdOptions::default();
36/// assert_eq!(options.derivative_eps, DEFAULT_DECOMPOSITION_DERIVATIVE_EPS);
37/// ```
38pub const DEFAULT_DECOMPOSITION_DERIVATIVE_EPS: f64 = 1e-12;
39
40/// Singular-vector gauge convention used by [`SvdOptions`].
41///
42/// # Examples
43///
44/// ```rust
45/// use tenferro_linalg::{SvdGauge, SvdOptions};
46///
47/// let options = SvdOptions::default().gauge(SvdGauge::CanonicalPivot);
48/// assert_eq!(options.gauge, SvdGauge::CanonicalPivot);
49/// ```
50#[derive(Clone, Copy, Debug, PartialEq, Eq)]
51pub enum SvdGauge {
52    /// Leave the backend's raw singular vector signs or phases unchanged.
53    Raw,
54    /// Make each left singular vector's max-absolute pivot entry positive-real
55    /// and adjust the matching `VT` row so reconstruction is preserved.
56    CanonicalPivot,
57}
58
59/// Eigenvector gauge convention used by [`EighOptions`].
60///
61/// # Examples
62///
63/// ```rust
64/// use tenferro_linalg::{EighGauge, EighOptions};
65///
66/// let options = EighOptions::default().gauge(EighGauge::CanonicalPivot);
67/// assert_eq!(options.gauge, EighGauge::CanonicalPivot);
68/// ```
69#[derive(Clone, Copy, Debug, PartialEq, Eq)]
70pub enum EighGauge {
71    /// Leave the backend's raw eigenvector signs or phases unchanged.
72    Raw,
73    /// Make each eigenvector's max-absolute pivot entry positive-real.
74    CanonicalPivot,
75}
76
77/// QR factor gauge convention used by [`QrOptions`].
78///
79/// # Examples
80///
81/// ```rust
82/// use tenferro_linalg::{QrGauge, QrOptions};
83///
84/// let options = QrOptions::default().gauge(QrGauge::PositiveDiagonal);
85/// assert_eq!(options.gauge, QrGauge::PositiveDiagonal);
86/// ```
87#[derive(Clone, Copy, Debug, PartialEq, Eq)]
88pub enum QrGauge {
89    /// Leave the backend's raw QR signs or phases unchanged.
90    Raw,
91    /// Make each `R` diagonal entry positive-real, compensating `Q`.
92    PositiveDiagonal,
93}
94
95/// Options for singular value decomposition.
96///
97/// # Examples
98///
99/// ```rust
100/// use tenferro_linalg::{SvdGauge, SvdOptions};
101///
102/// let options = SvdOptions::default()
103///     .gauge(SvdGauge::CanonicalPivot)
104///     .derivative_eps(1.0e-10);
105/// assert_eq!(options.gauge, SvdGauge::CanonicalPivot);
106/// assert_eq!(options.derivative_eps, 1.0e-10);
107/// ```
108#[derive(Clone, Copy, Debug, PartialEq)]
109pub struct SvdOptions {
110    /// Singular-vector gauge convention.
111    pub gauge: SvdGauge,
112    /// AD derivative regularization for repeated or nearly repeated singular values.
113    pub derivative_eps: f64,
114}
115
116impl Default for SvdOptions {
117    fn default() -> Self {
118        Self {
119            gauge: SvdGauge::Raw,
120            derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
121        }
122    }
123}
124
125impl SvdOptions {
126    /// Return options with the requested singular-vector gauge.
127    ///
128    /// # Examples
129    ///
130    /// ```rust
131    /// use tenferro_linalg::{SvdGauge, SvdOptions};
132    ///
133    /// let options = SvdOptions::default().gauge(SvdGauge::CanonicalPivot);
134    /// assert_eq!(options.gauge, SvdGauge::CanonicalPivot);
135    /// ```
136    pub fn gauge(mut self, gauge: SvdGauge) -> Self {
137        self.gauge = gauge;
138        self
139    }
140
141    /// Return options with an explicit derivative epsilon.
142    ///
143    /// # Examples
144    ///
145    /// ```rust
146    /// use tenferro_linalg::SvdOptions;
147    ///
148    /// let options = SvdOptions::default().derivative_eps(1.0e-9);
149    /// assert_eq!(options.derivative_eps, 1.0e-9);
150    /// ```
151    pub fn derivative_eps(mut self, derivative_eps: f64) -> Self {
152        self.derivative_eps = derivative_eps;
153        self
154    }
155}
156
157/// Options for Hermitian eigenvalue decomposition.
158///
159/// # Examples
160///
161/// ```rust
162/// use tenferro_linalg::EighOptions;
163///
164/// let options = EighOptions::default().derivative_eps(1.0e-10);
165/// assert_eq!(options.derivative_eps, 1.0e-10);
166/// ```
167#[derive(Clone, Copy, Debug, PartialEq)]
168pub struct EighOptions {
169    /// Eigenvector gauge convention.
170    pub gauge: EighGauge,
171    /// AD derivative regularization for repeated or nearly repeated eigenvalues.
172    pub derivative_eps: f64,
173}
174
175impl Default for EighOptions {
176    fn default() -> Self {
177        Self {
178            gauge: EighGauge::Raw,
179            derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
180        }
181    }
182}
183
184impl EighOptions {
185    /// Return options with the requested eigenvector gauge.
186    ///
187    /// # Examples
188    ///
189    /// ```rust
190    /// use tenferro_linalg::{EighGauge, EighOptions};
191    ///
192    /// let options = EighOptions::default().gauge(EighGauge::CanonicalPivot);
193    /// assert_eq!(options.gauge, EighGauge::CanonicalPivot);
194    /// ```
195    pub fn gauge(mut self, gauge: EighGauge) -> Self {
196        self.gauge = gauge;
197        self
198    }
199
200    /// Return options with an explicit derivative epsilon.
201    ///
202    /// # Examples
203    ///
204    /// ```rust
205    /// use tenferro_linalg::EighOptions;
206    ///
207    /// let options = EighOptions::default().derivative_eps(1.0e-9);
208    /// assert_eq!(options.derivative_eps, 1.0e-9);
209    /// ```
210    pub fn derivative_eps(mut self, derivative_eps: f64) -> Self {
211        self.derivative_eps = derivative_eps;
212        self
213    }
214}
215
216/// Options for QR decomposition.
217///
218/// # Examples
219///
220/// ```rust
221/// use tenferro_linalg::{QrGauge, QrOptions};
222///
223/// let options = QrOptions::default().gauge(QrGauge::PositiveDiagonal);
224/// assert_eq!(options.gauge, QrGauge::PositiveDiagonal);
225/// ```
226#[derive(Clone, Copy, Debug, PartialEq, Eq)]
227pub struct QrOptions {
228    /// QR sign or phase convention.
229    pub gauge: QrGauge,
230}
231
232impl Default for QrOptions {
233    fn default() -> Self {
234        Self {
235            gauge: QrGauge::Raw,
236        }
237    }
238}
239
240impl QrOptions {
241    /// Return options with the requested QR gauge.
242    ///
243    /// # Examples
244    ///
245    /// ```rust
246    /// use tenferro_linalg::{QrGauge, QrOptions};
247    ///
248    /// let options = QrOptions::default().gauge(QrGauge::PositiveDiagonal);
249    /// assert_eq!(options.gauge, QrGauge::PositiveDiagonal);
250    /// ```
251    pub fn gauge(mut self, gauge: QrGauge) -> Self {
252        self.gauge = gauge;
253        self
254    }
255}
256
257pub(crate) fn validate_derivative_eps(
258    op: &'static str,
259    derivative_eps: f64,
260) -> tenferro_tensor::Result<()> {
261    if derivative_eps.is_finite() && derivative_eps > 0.0 {
262        Ok(())
263    } else {
264        Err(Error::invalid_argument(
265            op,
266            "derivative_eps",
267            format!("must be positive and finite, got {derivative_eps}"),
268        ))
269    }
270}
271
272#[derive(Clone, Copy, Debug, PartialEq)]
273#[doc(hidden)]
274pub(crate) enum LinalgOp {
275    Cholesky,
276    Lu,
277    LuFactor,
278    LuSolvePrepared {
279        transpose_a: bool,
280        conjugate_a: bool,
281    },
282    SignDetFromLuFactor,
283    LogAbsDetFromLuFactor,
284    FullPivLu,
285    FullPivLuSolve {
286        transpose_a: bool,
287    },
288    /// Solve `a @ x = b` with partial-pivot LU (same kernel as
289    /// `LinalgBackend::solve`). Two inputs (matrix, rhs) to one output.
290    /// Only the eager surface (autodiff feature) constructs this variant;
291    /// the traced `solve` composite stays LuFactor + LuSolvePrepared.
292    #[cfg_attr(not(feature = "autodiff"), allow(dead_code))]
293    Solve,
294    Svd {
295        derivative_eps: f64,
296        gauge: SvdGauge,
297    },
298    /// Full-matrices SVD: `U` is `m x m` and `Vh` is `n x n`, so the trailing
299    /// `Vh` rows span the input's right nullspace. Value-only: AD is
300    /// intentionally unsupported (see the linalg AD support manifest).
301    SvdFull,
302    SvdVals {
303        derivative_eps: f64,
304    },
305    Qr {
306        gauge: QrGauge,
307    },
308    Eigh {
309        derivative_eps: f64,
310        gauge: EighGauge,
311    },
312    EighVals {
313        derivative_eps: f64,
314    },
315    Eig {
316        input_dtype: DType,
317    },
318    EigVals {
319        input_dtype: DType,
320    },
321    TriangularSolve {
322        left_side: bool,
323        lower: bool,
324        transpose_a: bool,
325        unit_diagonal: bool,
326    },
327}
328
329impl LinalgOp {
330    fn output_count(self) -> usize {
331        match self {
332            Self::Cholesky
333            | Self::EighVals { .. }
334            | Self::EigVals { .. }
335            | Self::FullPivLuSolve { .. }
336            | Self::LogAbsDetFromLuFactor
337            | Self::LuSolvePrepared { .. }
338            | Self::SignDetFromLuFactor
339            | Self::Solve
340            | Self::SvdVals { .. }
341            | Self::TriangularSolve { .. } => 1,
342            Self::Svd { .. } | Self::SvdFull => 3,
343            Self::Qr { .. } | Self::Eigh { .. } | Self::Eig { .. } => 2,
344            Self::LuFactor => 3,
345            Self::Lu => 4,
346            Self::FullPivLu => 5,
347        }
348    }
349
350    fn input_count(self) -> usize {
351        match self {
352            Self::FullPivLuSolve { .. } | Self::Solve | Self::TriangularSolve { .. } => 2,
353            Self::LogAbsDetFromLuFactor => 2,
354            Self::SignDetFromLuFactor => 3,
355            Self::LuSolvePrepared { .. } => 4,
356            _ => 1,
357        }
358    }
359
360    fn tag(self) -> u8 {
361        match self {
362            Self::Cholesky => 0,
363            Self::Lu => 1,
364            Self::FullPivLu => 2,
365            Self::FullPivLuSolve { .. } => 3,
366            Self::Svd { .. } => 4,
367            Self::Qr { .. } => 5,
368            Self::Eigh { .. } => 6,
369            Self::Eig { .. } => 7,
370            Self::TriangularSolve { .. } => 9,
371            Self::LuFactor => 10,
372            Self::LuSolvePrepared { .. } => 11,
373            Self::SvdVals { .. } => 12,
374            Self::EighVals { .. } => 13,
375            Self::EigVals { .. } => 14,
376            Self::SvdFull => 15,
377            Self::LogAbsDetFromLuFactor => 16,
378            Self::SignDetFromLuFactor => 17,
379            Self::Solve => 18,
380        }
381    }
382}
383
384#[derive(Clone, Debug, PartialEq)]
385#[doc(hidden)]
386pub(crate) struct LinalgExtensionOp {
387    op: LinalgOp,
388}
389
390impl LinalgExtensionOp {
391    pub(crate) fn new(op: LinalgOp) -> Self {
392        Self { op }
393    }
394
395    pub(crate) fn op(&self) -> LinalgOp {
396        self.op
397    }
398}
399
400impl ExtensionOp for LinalgExtensionOp {
401    fn family_id(&self) -> &'static str {
402        LINALG_EXTENSION_FAMILY_ID
403    }
404
405    fn payload_hash(&self, hasher: &mut dyn Hasher) {
406        hasher.write_u8(self.op.tag());
407        match self.op {
408            LinalgOp::Svd {
409                derivative_eps,
410                gauge,
411            } => {
412                hasher.write_u64(derivative_eps.to_bits());
413                hash_svd_gauge(hasher, gauge);
414            }
415            LinalgOp::SvdVals { derivative_eps } | LinalgOp::EighVals { derivative_eps } => {
416                hasher.write_u64(derivative_eps.to_bits());
417            }
418            LinalgOp::Qr { gauge } => {
419                hash_qr_gauge(hasher, gauge);
420            }
421            LinalgOp::Eigh {
422                derivative_eps,
423                gauge,
424            } => {
425                hasher.write_u64(derivative_eps.to_bits());
426                hash_eigh_gauge(hasher, gauge);
427            }
428            LinalgOp::Eig { input_dtype } | LinalgOp::EigVals { input_dtype } => {
429                hash_dtype(hasher, input_dtype);
430            }
431            LinalgOp::FullPivLuSolve { transpose_a } => {
432                hasher.write_u8(u8::from(transpose_a));
433            }
434            LinalgOp::LuSolvePrepared {
435                transpose_a,
436                conjugate_a,
437            } => {
438                hasher.write_u8(u8::from(transpose_a));
439                hasher.write_u8(u8::from(conjugate_a));
440            }
441            LinalgOp::TriangularSolve {
442                left_side,
443                lower,
444                transpose_a,
445                unit_diagonal,
446            } => {
447                hasher.write_u8(u8::from(left_side));
448                hasher.write_u8(u8::from(lower));
449                hasher.write_u8(u8::from(transpose_a));
450                hasher.write_u8(u8::from(unit_diagonal));
451            }
452            LinalgOp::Cholesky
453            | LinalgOp::Lu
454            | LinalgOp::LuFactor
455            | LinalgOp::LogAbsDetFromLuFactor
456            | LinalgOp::SignDetFromLuFactor
457            | LinalgOp::FullPivLu
458            | LinalgOp::SvdFull
459            | LinalgOp::Solve => {}
460        }
461    }
462
463    fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
464        other
465            .as_any()
466            .downcast_ref::<Self>()
467            .is_some_and(|that| self == that)
468    }
469
470    fn clone_arc(&self) -> Arc<dyn ExtensionOp> {
471        Arc::new(self.clone())
472    }
473
474    fn as_any(&self) -> &dyn Any {
475        self
476    }
477
478    fn input_count(&self) -> usize {
479        self.op.input_count()
480    }
481
482    fn output_count(&self) -> usize {
483        self.op.output_count()
484    }
485
486    fn semantic_effects(&self) -> tenferro_ops::ext_op::ExtensionEffectDeclaration<'_> {
487        tenferro_ops::ext_op::ExtensionEffectDeclaration::Declared(&[])
488    }
489
490    fn semantic_aliases(&self) -> tenferro_ops::ext_op::ExtensionAliasDeclaration<'_> {
491        tenferro_ops::ext_op::ExtensionAliasDeclaration::AllFresh
492    }
493
494    fn prune_outputs(&self, live_outputs: &[bool]) -> Option<Arc<dyn ExtensionOp>> {
495        match self.op {
496            LinalgOp::Svd { derivative_eps, .. } if live_outputs == [false, true, false] => {
497                Some(Arc::new(Self::new(LinalgOp::SvdVals { derivative_eps })))
498            }
499            LinalgOp::Eigh { derivative_eps, .. } if live_outputs == [true, false] => {
500                Some(Arc::new(Self::new(LinalgOp::EighVals { derivative_eps })))
501            }
502            LinalgOp::Eig { input_dtype } if live_outputs == [true, false] => {
503                Some(Arc::new(Self::new(LinalgOp::EigVals { input_dtype })))
504            }
505            _ => None,
506        }
507    }
508
509    fn infer_output_meta(
510        &self,
511        ctx: &mut tenferro_ops::ExtensionShapeContext<'_>,
512    ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
513        let input_dtypes = (0..self.input_count())
514            .map(|input| ctx.input_dtype(input))
515            .collect::<Result<Vec<_>, _>>()?;
516        let input_shapes = (0..self.input_count())
517            .map(|input| ctx.input_shape(input))
518            .collect::<Result<Vec<_>, _>>()?;
519        let metas = match self.op {
520            LinalgOp::Cholesky => {
521                require_matrix_meta("tenferro-linalg.cholesky", input_shapes[0])?;
522                vec![(promote_dtypes(&input_dtypes), input_shapes[0].to_vec())]
523            }
524            LinalgOp::FullPivLuSolve { .. } => {
525                require_matrix_meta("tenferro-linalg.full_piv_lu_solve", input_shapes[0])?;
526                require_matrix_meta("tenferro-linalg.full_piv_lu_solve", input_shapes[1])?;
527                vec![(promote_dtypes(&input_dtypes), input_shapes[1].to_vec())]
528            }
529            LinalgOp::Solve => {
530                require_matrix_meta("tenferro-linalg.solve", input_shapes[0])?;
531                require_matrix_meta("tenferro-linalg.solve", input_shapes[1])?;
532                vec![(promote_dtypes(&input_dtypes), input_shapes[1].to_vec())]
533            }
534            LinalgOp::TriangularSolve { .. } => {
535                require_matrix_meta("tenferro-linalg.triangular_solve", input_shapes[0])?;
536                require_matrix_meta("tenferro-linalg.triangular_solve", input_shapes[1])?;
537                vec![(promote_dtypes(&input_dtypes), input_shapes[1].to_vec())]
538            }
539            LinalgOp::LuSolvePrepared { .. } => {
540                require_matrix_meta("tenferro-linalg.lu_solve_prepared_lu", input_shapes[0])?;
541                require_matrix_meta("tenferro-linalg.lu_solve_prepared_rhs", input_shapes[3])?;
542                vec![(
543                    promote_dtypes(&[input_dtypes[0], input_dtypes[3]]),
544                    input_shapes[3].to_vec(),
545                )]
546            }
547            LinalgOp::Lu => lu_meta(input_dtypes[0], input_shapes[0])?,
548            LinalgOp::LuFactor => lu_factor_meta(input_dtypes[0], input_shapes[0])?,
549            LinalgOp::SignDetFromLuFactor => {
550                vec![signdet_from_lu_factor_meta(
551                    input_dtypes[0],
552                    input_shapes[0],
553                    input_shapes[1],
554                    input_shapes[2],
555                )?]
556            }
557            LinalgOp::LogAbsDetFromLuFactor => {
558                vec![logabsdet_from_lu_factor_meta(
559                    input_dtypes[0],
560                    input_shapes[0],
561                    input_shapes[1],
562                )?]
563            }
564            LinalgOp::FullPivLu => full_piv_lu_meta(input_dtypes[0], input_shapes[0])?,
565            LinalgOp::Svd { .. } => svd_meta(input_dtypes[0], input_shapes[0])?,
566            LinalgOp::SvdFull => svd_full_meta(input_dtypes[0], input_shapes[0])?,
567            LinalgOp::SvdVals { .. } => {
568                vec![svd_values_meta(input_dtypes[0], input_shapes[0])?]
569            }
570            LinalgOp::Qr { .. } => qr_meta(input_dtypes[0], input_shapes[0])?,
571            LinalgOp::Eigh { .. } => eigh_meta(input_dtypes[0], input_shapes[0])?,
572            LinalgOp::EighVals { .. } => vec![eigh_values_meta(input_dtypes[0], input_shapes[0])?],
573            LinalgOp::Eig { input_dtype } => eig_meta(input_dtype, input_shapes[0])?,
574            LinalgOp::EigVals { input_dtype } => {
575                vec![eig_values_meta(input_dtype, input_shapes[0])?]
576            }
577        };
578        Ok(metas)
579    }
580}
581
582pub(crate) fn execute_linalg_extension_reads<B: BackendSession + ?Sized>(
583    op: &LinalgExtensionOp,
584    inputs: &[TensorRead<'_>],
585    ctx: &mut ExtensionExecutionContext<'_, B>,
586) -> tenferro_tensor::Result<Vec<Tensor>> {
587    execute_linalg_extension_reads_on_session(op, inputs, ctx.backend_mut())
588}
589
590pub(crate) fn execute_linalg_extension_reads_owner<B: TensorBackend>(
591    op: &LinalgExtensionOp,
592    inputs: &[TensorRead<'_>],
593    ctx: &mut ExtensionExecutionContext<'_, B>,
594) -> tenferro_tensor::Result<Vec<Tensor>> {
595    let (backend, caches) = ctx.parts_mut();
596    backend.with_backend_session(|session| {
597        let mut session_ctx = ExtensionExecutionContext::new(session, caches);
598        execute_linalg_extension_reads(op, inputs, &mut session_ctx)
599    })
600}
601
602fn execute_linalg_extension_reads_on_session<B: BackendSession + ?Sized>(
603    op: &LinalgExtensionOp,
604    inputs: &[TensorRead<'_>],
605    session: &mut B,
606) -> tenferro_tensor::Result<Vec<Tensor>> {
607    if let Some(result) = with_cpu_exec_session(session, |session| {
608        execute_linalg_extension_reads_in_session(op, inputs, session)
609    }) {
610        return result;
611    }
612    #[cfg(feature = "cuda")]
613    if let Some(result) = with_cuda_exec_session(session, |session| {
614        execute_linalg_extension_reads_in_session(op, inputs, session)
615    }) {
616        return result;
617    }
618    Err(Error::unsupported(
619        "linalg_extension",
620        "selected backend session does not expose a linalg execution capability",
621    ))
622}
623
624fn execute_linalg_extension_reads_in_session<S: LinalgBackend>(
625    op: &LinalgExtensionOp,
626    inputs: &[TensorRead<'_>],
627    session: &mut S,
628) -> tenferro_tensor::Result<Vec<Tensor>> {
629    if op.op() == LinalgOp::Cholesky {
630        return Ok(vec![session.cholesky_read(inputs[0].clone())?]);
631    }
632    if let LinalgOp::TriangularSolve {
633        left_side,
634        lower,
635        transpose_a,
636        unit_diagonal,
637    } = op.op()
638    {
639        match session.triangular_solve_read(
640            inputs[0].clone(),
641            inputs[1].clone(),
642            left_side,
643            lower,
644            transpose_a,
645            unit_diagonal,
646        ) {
647            Ok(output) => return Ok(vec![output]),
648            Err(error) if error.kind() == ErrorKind::Unsupported => {}
649            Err(error) => return Err(error),
650        }
651    }
652
653    // Linalg kernels operate on compact tensors; materialization is explicit
654    // here so borrowed views cannot bypass provider errors.
655    let materialized_inputs = inputs
656        .iter()
657        .cloned()
658        .map(|input| session.to_contiguous_read(input))
659        .collect::<tenferro_tensor::Result<Vec<_>>>()?;
660    let input_refs: Vec<&Tensor> = materialized_inputs.iter().collect();
661    execute_linalg(op.op(), &input_refs, session)
662}
663
664fn linalg_session_supported<B: BackendSession + 'static>(op: &LinalgExtensionOp) -> bool {
665    // The `supports_session` contract (capability.rs) requires that an op is
666    // admitted to a scheduler session only when the session executor genuinely
667    // executes it without returning `Unsupported`. Admission is exactly
668    // per-op/per-backend so `apply_eager` keeps the native prepared path for
669    // every op the session can actually run (issue #1665).
670    let type_id = std::any::TypeId::of::<B>();
671    if type_id == std::any::TypeId::of::<tenferro_cpu::CpuBackend>() {
672        // The CPU backend type does not carry its provider kind (faer vs BLAS)
673        // at this type-only seam, and the BLAS provider does not implement
674        // in-session full-matrices SVD, so SvdFull is conservatively rejected
675        // and falls back to the compiled path. Every other CPU linalg kernel
676        // runs in-session on both faer and BLAS providers.
677        return op.op() != LinalgOp::SvdFull;
678    }
679    #[cfg(feature = "cuda")]
680    {
681        if type_id == std::any::TypeId::of::<tenferro_gpu::cuda::CudaBackend>() {
682            return match op.op() {
683                // Complete-pivoting LU and general eig have no CUDA kernels.
684                LinalgOp::FullPivLu | LinalgOp::FullPivLuSolve { .. } => false,
685                LinalgOp::Eig { .. } | LinalgOp::EigVals { .. } => false,
686                // Plain partial-pivot solve runs in-session via cuSOLVER
687                // getrf plus prepared pivot/triangular solves
688                // (`gpu/linalg.rs::solve` = lu_factor + lu_solve_prepared, no
689                // Unsupported path for F32/F64/C32/C64), so it is admitted.
690                LinalgOp::Solve => true,
691                // Full-matrices SVD falls back to the default `svd_full`
692                // impl, which reports `Unsupported`.
693                LinalgOp::SvdFull => false,
694                // Conjugate-only prepared LU solve is unsupported on CUDA.
695                LinalgOp::LuSolvePrepared {
696                    transpose_a: false,
697                    conjugate_a: true,
698                } => false,
699                _ => true,
700            };
701        }
702    }
703    false
704}
705
706fn execute_linalg_extension_in_session(
707    op: &LinalgExtensionOp,
708    session: &mut dyn BackendSession,
709    _extension_caches: &mut tenferro_runtime::ExtensionCacheStore,
710    inputs: &[TensorRead<'_>],
711) -> tenferro_tensor::Result<Vec<Tensor>> {
712    // Reuse the existing session executor that the eager and scheduler paths
713    // already share; it downcasts the borrowed session to the CPU/CUDA exec
714    // session and runs the same forward kernel for every LinalgOp.
715    execute_linalg_extension_reads_on_session(op, inputs, session)
716}
717
718define_extension_runtime! {
719    runtime = LinalgRuntime,
720    family_id = LINALG_EXTENSION_FAMILY_ID,
721    op_type = LinalgExtensionOp,
722    execute = execute_linalg_extension_reads_owner,
723    execute_reads = execute_linalg_extension_reads_owner,
724    execute_in_session = execute_linalg_extension_in_session,
725    session_supported = linalg_session_supported,
726    backend_bound = TensorBackend,
727}
728
729fn execute_linalg<B: LinalgBackend>(
730    op: LinalgOp,
731    inputs: &[&Tensor],
732    backend: &mut B,
733) -> tenferro_tensor::Result<Vec<Tensor>> {
734    match op {
735        LinalgOp::Cholesky => Ok(vec![backend.cholesky(inputs[0])?]),
736        LinalgOp::Lu => backend.lu(inputs[0]),
737        LinalgOp::LuFactor => backend.lu_factor(inputs[0]),
738        LinalgOp::SignDetFromLuFactor => Ok(vec![signdet_from_lu_factor(
739            inputs[0].dtype(),
740            inputs[1],
741            inputs[2],
742            backend,
743        )?]),
744        LinalgOp::LogAbsDetFromLuFactor => Ok(vec![logabsdet_from_lu_factor(inputs[1], backend)?]),
745        LinalgOp::LuSolvePrepared {
746            transpose_a,
747            conjugate_a,
748        } => Ok(vec![backend.lu_solve_prepared(
749            inputs[0],
750            inputs[1],
751            inputs[2],
752            inputs[3],
753            transpose_a,
754            conjugate_a,
755        )?]),
756        LinalgOp::FullPivLu => backend.full_piv_lu(inputs[0]),
757        LinalgOp::FullPivLuSolve { transpose_a } => Ok(vec![backend.full_piv_lu_solve(
758            inputs[0],
759            inputs[1],
760            transpose_a,
761        )?]),
762        LinalgOp::Solve => Ok(vec![backend.solve(inputs[0], inputs[1])?]),
763        LinalgOp::Svd {
764            derivative_eps,
765            gauge,
766        } => backend.svd_with_options(
767            inputs[0],
768            SvdOptions {
769                derivative_eps,
770                gauge,
771            },
772        ),
773        LinalgOp::SvdFull => backend.svd_full(inputs[0]),
774        LinalgOp::SvdVals { .. } => Ok(vec![backend.svd_values(inputs[0])?]),
775        LinalgOp::Qr { gauge } => backend.qr_with_options(inputs[0], QrOptions { gauge }),
776        LinalgOp::Eigh {
777            derivative_eps,
778            gauge,
779        } => backend.eigh_with_options(
780            inputs[0],
781            EighOptions {
782                derivative_eps,
783                gauge,
784            },
785        ),
786        LinalgOp::EighVals { .. } => Ok(vec![backend.eigh_values(inputs[0])?]),
787        LinalgOp::Eig { .. } => backend.eig(inputs[0]),
788        LinalgOp::EigVals { .. } => Ok(vec![backend.eig_values(inputs[0])?]),
789        LinalgOp::TriangularSolve {
790            left_side,
791            lower,
792            transpose_a,
793            unit_diagonal,
794        } => Ok(vec![backend.triangular_solve(
795            inputs[0],
796            inputs[1],
797            left_side,
798            lower,
799            transpose_a,
800            unit_diagonal,
801        )?]),
802    }
803}
804
805fn signdet_from_lu_factor<B: LinalgBackend + ?Sized>(
806    input_dtype: DType,
807    packed_lu: &Tensor,
808    parity: &Tensor,
809    backend: &mut B,
810) -> tenferro_tensor::Result<Tensor> {
811    let diag = backend.extract_diagonal(packed_lu, 0, 1)?;
812    let det_u = backend.reduce_prod_read(TensorRead::from_tensor(&diag), &[0])?;
813    let det = backend.mul_read(
814        TensorRead::from_tensor(parity),
815        TensorRead::from_tensor(&det_u),
816    )?;
817    if matches!(input_dtype, DType::C32 | DType::C64) {
818        let abs = backend.abs_read(TensorRead::from_tensor(&det))?;
819        let abs = backend.convert(&abs, input_dtype)?;
820        backend.div_read(TensorRead::from_tensor(&det), TensorRead::from_tensor(&abs))
821    } else {
822        backend.sign_read(TensorRead::from_tensor(&det))
823    }
824}
825
826fn logabsdet_from_lu_factor<B: LinalgBackend + ?Sized>(
827    packed_lu: &Tensor,
828    backend: &mut B,
829) -> tenferro_tensor::Result<Tensor> {
830    let diag = backend.extract_diagonal(packed_lu, 0, 1)?;
831    let abs = backend.abs_read(TensorRead::from_tensor(&diag))?;
832    let log = backend.log_read(TensorRead::from_tensor(&abs))?;
833    backend.reduce_sum_read(TensorRead::from_tensor(&log), &[0])
834}
835
836pub(crate) fn apply_svd_gauge(
837    gauge: SvdGauge,
838    outputs: &mut [Tensor],
839) -> tenferro_tensor::Result<()> {
840    match gauge {
841        SvdGauge::Raw => Ok(()),
842        SvdGauge::CanonicalPivot => apply_canonical_pivot_svd_gauge(outputs),
843    }
844}
845
846fn apply_canonical_pivot_svd_gauge(outputs: &mut [Tensor]) -> tenferro_tensor::Result<()> {
847    if outputs.len() != 3 {
848        return Err(Error::invalid_argument(
849            "tenferro-linalg.svd",
850            "outputs",
851            format!(
852                "canonical SVD gauge expected three outputs, got {}",
853                outputs.len()
854            ),
855        ));
856    }
857
858    let (u_slice, rest) = outputs.split_at_mut(1);
859    let (singular_slice, vt_slice) = rest.split_at_mut(1);
860    let u = &mut u_slice[0];
861    let singular_values = &singular_slice[0];
862    let vt = &mut vt_slice[0];
863    let u_shape = u.shape().to_vec();
864    let s_shape = singular_values.shape().to_vec();
865    let vt_shape = vt.shape().to_vec();
866    if u_shape.len() < 2 || vt_shape.len() < 2 || s_shape.is_empty() {
867        return Err(Error::invalid_argument(
868            "tenferro-linalg.svd",
869            "outputs",
870            format!(
871                "canonical SVD gauge expected U rank >= 2, S rank >= 1, VT rank >= 2; got U={u_shape:?}, S={s_shape:?}, VT={vt_shape:?}"
872            ),
873        ));
874    }
875
876    let m = u_shape[0];
877    let k = u_shape[1];
878    let n = vt_shape[1];
879    if s_shape[0] != k
880        || vt_shape[0] != k
881        || u_shape[2..] != vt_shape[2..]
882        || s_shape[1..] != u_shape[2..]
883    {
884        return Err(Error::invalid_argument(
885            "tenferro-linalg.svd",
886            "outputs",
887            format!(
888                "canonical SVD gauge expected compatible compact SVD shapes, got U={u_shape:?}, S={s_shape:?}, VT={vt_shape:?}"
889            ),
890        ));
891    }
892    let layout = canonical_svd_gauge_layout(m, k, n, &u_shape[2..])?;
893
894    match (u, vt) {
895        (Tensor::F64(u), Tensor::F64(vt)) => {
896            canonicalize_svd_gauge_f64(u.host_data_mut()?, vt.host_data_mut()?, layout)
897        }
898        (Tensor::F32(u), Tensor::F32(vt)) => {
899            canonicalize_svd_gauge_f32(u.host_data_mut()?, vt.host_data_mut()?, layout)
900        }
901        (Tensor::C64(u), Tensor::C64(vt)) => {
902            canonicalize_svd_gauge_c64(u.host_data_mut()?, vt.host_data_mut()?, layout)
903        }
904        (Tensor::C32(u), Tensor::C32(vt)) => {
905            canonicalize_svd_gauge_c32(u.host_data_mut()?, vt.host_data_mut()?, layout)
906        }
907        (u, vt) => Err(Error::dtype_mismatch(
908            "tenferro-linalg.svd",
909            u.dtype(),
910            vt.dtype(),
911        )),
912    }
913}
914
915#[derive(Clone, Copy, Debug, PartialEq, Eq)]
916struct CanonicalSvdGaugeLayout {
917    m: usize,
918    k: usize,
919    batch_count: usize,
920    u_batch_len: usize,
921    vt_batch_len: usize,
922    u_len: usize,
923    vt_len: usize,
924}
925
926impl CanonicalSvdGaugeLayout {
927    fn validate_storage(self, u_len: usize, vt_len: usize) -> tenferro_tensor::Result<()> {
928        if u_len != self.u_len {
929            return Err(Error::invalid_argument(
930                "tenferro-linalg.svd",
931                "U storage",
932                format!(
933                    "canonical SVD gauge expected U storage length {}, got {u_len}",
934                    self.u_len
935                ),
936            ));
937        }
938        if vt_len != self.vt_len {
939            return Err(Error::invalid_argument(
940                "tenferro-linalg.svd",
941                "VT storage",
942                format!(
943                    "canonical SVD gauge expected VT storage length {}, got {vt_len}",
944                    self.vt_len
945                ),
946            ));
947        }
948        Ok(())
949    }
950}
951
952fn canonical_svd_gauge_layout(
953    m: usize,
954    k: usize,
955    n: usize,
956    batch_shape: &[usize],
957) -> tenferro_tensor::Result<CanonicalSvdGaugeLayout> {
958    let batch_count = tenferro_tensor::validate::checked_shape_product(
959        "tenferro-linalg.svd",
960        "canonical SVD batch",
961        batch_shape,
962    )?;
963    let u_batch_len = tenferro_tensor::validate::checked_shape_product(
964        "tenferro-linalg.svd",
965        "canonical SVD U batch",
966        &[m, k],
967    )?;
968    let vt_batch_len = tenferro_tensor::validate::checked_shape_product(
969        "tenferro-linalg.svd",
970        "canonical SVD VT batch",
971        &[k, n],
972    )?;
973    let u_len = tenferro_tensor::validate::checked_shape_product(
974        "tenferro-linalg.svd",
975        "canonical SVD U storage",
976        &[u_batch_len, batch_count],
977    )?;
978    let vt_len = tenferro_tensor::validate::checked_shape_product(
979        "tenferro-linalg.svd",
980        "canonical SVD VT storage",
981        &[vt_batch_len, batch_count],
982    )?;
983    Ok(CanonicalSvdGaugeLayout {
984        m,
985        k,
986        batch_count,
987        u_batch_len,
988        vt_batch_len,
989        u_len,
990        vt_len,
991    })
992}
993
994fn canonicalize_svd_gauge_f64(
995    u: &mut [f64],
996    vt: &mut [f64],
997    layout: CanonicalSvdGaugeLayout,
998) -> tenferro_tensor::Result<()> {
999    layout.validate_storage(u.len(), vt.len())?;
1000    if layout.batch_count == 0 || layout.u_batch_len == 0 || layout.vt_batch_len == 0 {
1001        return Ok(());
1002    }
1003    for (u_batch, vt_batch) in u
1004        .chunks_exact_mut(layout.u_batch_len)
1005        .zip(vt.chunks_exact_mut(layout.vt_batch_len))
1006    {
1007        for (col, u_column) in u_batch.chunks_exact_mut(layout.m).enumerate() {
1008            let pivot = max_abs_pivot_f64(u_column);
1009            let pivot_value = u_column[pivot];
1010            if pivot_value < 0.0 {
1011                for value in u_column {
1012                    *value = -*value;
1013                }
1014                for vt_column in vt_batch.chunks_exact_mut(layout.k) {
1015                    vt_column[col] = -vt_column[col];
1016                }
1017            }
1018        }
1019    }
1020    Ok(())
1021}
1022
1023fn canonicalize_svd_gauge_f32(
1024    u: &mut [f32],
1025    vt: &mut [f32],
1026    layout: CanonicalSvdGaugeLayout,
1027) -> tenferro_tensor::Result<()> {
1028    layout.validate_storage(u.len(), vt.len())?;
1029    if layout.batch_count == 0 || layout.u_batch_len == 0 || layout.vt_batch_len == 0 {
1030        return Ok(());
1031    }
1032    for (u_batch, vt_batch) in u
1033        .chunks_exact_mut(layout.u_batch_len)
1034        .zip(vt.chunks_exact_mut(layout.vt_batch_len))
1035    {
1036        for (col, u_column) in u_batch.chunks_exact_mut(layout.m).enumerate() {
1037            let pivot = max_abs_pivot_f32(u_column);
1038            let pivot_value = u_column[pivot];
1039            if pivot_value < 0.0 {
1040                for value in u_column {
1041                    *value = -*value;
1042                }
1043                for vt_column in vt_batch.chunks_exact_mut(layout.k) {
1044                    vt_column[col] = -vt_column[col];
1045                }
1046            }
1047        }
1048    }
1049    Ok(())
1050}
1051
1052fn canonicalize_svd_gauge_c64(
1053    u: &mut [Complex64],
1054    vt: &mut [Complex64],
1055    layout: CanonicalSvdGaugeLayout,
1056) -> tenferro_tensor::Result<()> {
1057    layout.validate_storage(u.len(), vt.len())?;
1058    if layout.batch_count == 0 || layout.u_batch_len == 0 || layout.vt_batch_len == 0 {
1059        return Ok(());
1060    }
1061    for (u_batch, vt_batch) in u
1062        .chunks_exact_mut(layout.u_batch_len)
1063        .zip(vt.chunks_exact_mut(layout.vt_batch_len))
1064    {
1065        for (col, u_column) in u_batch.chunks_exact_mut(layout.m).enumerate() {
1066            let pivot = max_abs_pivot_c64(u_column);
1067            let pivot_value = u_column[pivot];
1068            let pivot_norm = pivot_value.norm();
1069            if pivot_norm == 0.0 {
1070                continue;
1071            }
1072            let phase = pivot_value.conj() / pivot_norm;
1073            let vt_phase = phase.conj();
1074            for value in u_column {
1075                *value *= phase;
1076            }
1077            for vt_column in vt_batch.chunks_exact_mut(layout.k) {
1078                vt_column[col] *= vt_phase;
1079            }
1080        }
1081    }
1082    Ok(())
1083}
1084
1085fn canonicalize_svd_gauge_c32(
1086    u: &mut [Complex32],
1087    vt: &mut [Complex32],
1088    layout: CanonicalSvdGaugeLayout,
1089) -> tenferro_tensor::Result<()> {
1090    layout.validate_storage(u.len(), vt.len())?;
1091    if layout.batch_count == 0 || layout.u_batch_len == 0 || layout.vt_batch_len == 0 {
1092        return Ok(());
1093    }
1094    for (u_batch, vt_batch) in u
1095        .chunks_exact_mut(layout.u_batch_len)
1096        .zip(vt.chunks_exact_mut(layout.vt_batch_len))
1097    {
1098        for (col, u_column) in u_batch.chunks_exact_mut(layout.m).enumerate() {
1099            let pivot = max_abs_pivot_c32(u_column);
1100            let pivot_value = u_column[pivot];
1101            let pivot_norm = pivot_value.norm();
1102            if pivot_norm == 0.0 {
1103                continue;
1104            }
1105            let phase = pivot_value.conj() / pivot_norm;
1106            let vt_phase = phase.conj();
1107            for value in u_column {
1108                *value *= phase;
1109            }
1110            for vt_column in vt_batch.chunks_exact_mut(layout.k) {
1111                vt_column[col] *= vt_phase;
1112            }
1113        }
1114    }
1115    Ok(())
1116}
1117
1118fn max_abs_pivot_f64(u_column: &[f64]) -> usize {
1119    let mut pivot = 0;
1120    let mut pivot_abs = u_column[0].abs();
1121    for (row, value) in u_column.iter().enumerate().skip(1) {
1122        let candidate_abs = value.abs();
1123        if candidate_abs > pivot_abs {
1124            pivot = row;
1125            pivot_abs = candidate_abs;
1126        }
1127    }
1128    pivot
1129}
1130
1131fn max_abs_pivot_f32(u_column: &[f32]) -> usize {
1132    let mut pivot = 0;
1133    let mut pivot_abs = u_column[0].abs();
1134    for (row, value) in u_column.iter().enumerate().skip(1) {
1135        let candidate_abs = value.abs();
1136        if candidate_abs > pivot_abs {
1137            pivot = row;
1138            pivot_abs = candidate_abs;
1139        }
1140    }
1141    pivot
1142}
1143
1144fn max_abs_pivot_c64(u_column: &[Complex64]) -> usize {
1145    let mut pivot = 0;
1146    let mut pivot_abs = u_column[0].norm_sqr();
1147    for (row, value) in u_column.iter().enumerate().skip(1) {
1148        let candidate_abs = value.norm_sqr();
1149        if candidate_abs > pivot_abs {
1150            pivot = row;
1151            pivot_abs = candidate_abs;
1152        }
1153    }
1154    pivot
1155}
1156
1157fn max_abs_pivot_c32(u_column: &[Complex32]) -> usize {
1158    let mut pivot = 0;
1159    let mut pivot_abs = u_column[0].norm_sqr();
1160    for (row, value) in u_column.iter().enumerate().skip(1) {
1161        let candidate_abs = value.norm_sqr();
1162        if candidate_abs > pivot_abs {
1163            pivot = row;
1164            pivot_abs = candidate_abs;
1165        }
1166    }
1167    pivot
1168}
1169
1170fn require_matrix_meta(op: &'static str, shape: &[SymDim]) -> tenferro_tensor::Result<()> {
1171    if shape.len() < 2 {
1172        return Err(Error::rank_mismatch(op, 2, shape.len()));
1173    }
1174    Ok(())
1175}
1176
1177fn matrix_meta_parts<'a>(
1178    op: &'static str,
1179    shape: &'a [SymDim],
1180) -> tenferro_tensor::Result<(SymDim, SymDim, &'a [SymDim])> {
1181    require_matrix_meta(op, shape)?;
1182    Ok((shape[0].clone(), shape[1].clone(), &shape[2..]))
1183}
1184
1185fn lu_meta(dtype: DType, shape: &[SymDim]) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1186    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.lu", shape)?;
1187    let k = m.clone().min(n.clone());
1188    Ok(vec![
1189        (dtype, matrix_shape(m.clone(), m, batch)),
1190        (dtype, matrix_shape(shape[0].clone(), k.clone(), batch)),
1191        (dtype, matrix_shape(k, n, batch)),
1192        (dtype, batch.to_vec()),
1193    ])
1194}
1195
1196fn lu_factor_meta(
1197    dtype: DType,
1198    shape: &[SymDim],
1199) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1200    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.lu_factor", shape)?;
1201    let k = m.min(n);
1202    Ok(vec![
1203        (dtype, shape.to_vec()),
1204        (DType::I32, vector_shape(k, batch)),
1205        (dtype, batch.to_vec()),
1206    ])
1207}
1208
1209fn signdet_from_lu_factor_meta(
1210    input_dtype: DType,
1211    input_shape: &[SymDim],
1212    packed_shape: &[SymDim],
1213    parity_shape: &[SymDim],
1214) -> tenferro_tensor::Result<(DType, Vec<SymDim>)> {
1215    let (_, _, batch) = matrix_meta_parts("tenferro-linalg.signdet_from_lu_factor", input_shape)?;
1216    require_matrix_meta(
1217        "tenferro-linalg.signdet_from_lu_factor_packed",
1218        packed_shape,
1219    )?;
1220    if parity_shape.len() != batch.len() {
1221        return Err(Error::rank_mismatch(
1222            "tenferro-linalg.signdet_from_lu_factor_parity",
1223            batch.len(),
1224            parity_shape.len(),
1225        ));
1226    }
1227    Ok((input_dtype, batch.to_vec()))
1228}
1229
1230fn logabsdet_from_lu_factor_meta(
1231    input_dtype: DType,
1232    input_shape: &[SymDim],
1233    packed_shape: &[SymDim],
1234) -> tenferro_tensor::Result<(DType, Vec<SymDim>)> {
1235    let (_, _, batch) = matrix_meta_parts("tenferro-linalg.logabsdet_from_lu_factor", input_shape)?;
1236    require_matrix_meta(
1237        "tenferro-linalg.logabsdet_from_lu_factor_packed",
1238        packed_shape,
1239    )?;
1240    Ok((singular_values_dtype(input_dtype), batch.to_vec()))
1241}
1242
1243fn full_piv_lu_meta(
1244    dtype: DType,
1245    shape: &[SymDim],
1246) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1247    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.full_piv_lu", shape)?;
1248    Ok(vec![
1249        (dtype, matrix_shape(n.clone(), n.clone(), batch)),
1250        (dtype, matrix_shape(n.clone(), n.clone(), batch)),
1251        (dtype, matrix_shape(n.clone(), n.clone(), batch)),
1252        (dtype, matrix_shape(n.clone(), n, batch)),
1253        (singular_values_dtype(dtype), batch.to_vec()),
1254    ])
1255}
1256
1257fn svd_meta(dtype: DType, shape: &[SymDim]) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1258    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.svd", shape)?;
1259    let k = m.clone().min(n.clone());
1260    Ok(vec![
1261        (dtype, matrix_shape(m, k.clone(), batch)),
1262        (singular_values_dtype(dtype), vector_shape(k.clone(), batch)),
1263        (dtype, matrix_shape(k, n, batch)),
1264    ])
1265}
1266
1267fn svd_full_meta(
1268    dtype: DType,
1269    shape: &[SymDim],
1270) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1271    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.svd_full", shape)?;
1272    let k = m.clone().min(n.clone());
1273    Ok(vec![
1274        (dtype, matrix_shape(m.clone(), m, batch)),
1275        (singular_values_dtype(dtype), vector_shape(k, batch)),
1276        (dtype, matrix_shape(n.clone(), n, batch)),
1277    ])
1278}
1279
1280fn svd_values_meta(
1281    dtype: DType,
1282    shape: &[SymDim],
1283) -> tenferro_tensor::Result<(DType, Vec<SymDim>)> {
1284    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.svd_values", shape)?;
1285    let k = m.min(n);
1286    Ok((singular_values_dtype(dtype), vector_shape(k, batch)))
1287}
1288
1289fn qr_meta(dtype: DType, shape: &[SymDim]) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1290    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.qr", shape)?;
1291    let k = m.clone().min(n.clone());
1292    Ok(vec![
1293        (dtype, matrix_shape(m, k.clone(), batch)),
1294        (dtype, matrix_shape(k, n, batch)),
1295    ])
1296}
1297
1298fn eigh_meta(dtype: DType, shape: &[SymDim]) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1299    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.eigh", shape)?;
1300    Ok(vec![
1301        (singular_values_dtype(dtype), vector_shape(n.clone(), batch)),
1302        (dtype, matrix_shape(n.clone(), n, batch)),
1303    ])
1304}
1305
1306fn eigh_values_meta(
1307    dtype: DType,
1308    shape: &[SymDim],
1309) -> tenferro_tensor::Result<(DType, Vec<SymDim>)> {
1310    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.eigh_values", shape)?;
1311    Ok((singular_values_dtype(dtype), vector_shape(n, batch)))
1312}
1313
1314fn eig_meta(
1315    input_dtype: DType,
1316    shape: &[SymDim],
1317) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1318    let dtype = eig_output_dtype(input_dtype);
1319    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.eig", shape)?;
1320    Ok(vec![
1321        (dtype, vector_shape(n.clone(), batch)),
1322        (dtype, matrix_shape(n.clone(), n, batch)),
1323    ])
1324}
1325
1326fn eig_values_meta(
1327    input_dtype: DType,
1328    shape: &[SymDim],
1329) -> tenferro_tensor::Result<(DType, Vec<SymDim>)> {
1330    let dtype = eig_output_dtype(input_dtype);
1331    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.eig_values", shape)?;
1332    Ok((dtype, vector_shape(n, batch)))
1333}
1334
1335fn matrix_shape(rows: SymDim, cols: SymDim, batch: &[SymDim]) -> Vec<SymDim> {
1336    let mut shape = vec![rows, cols];
1337    shape.extend_from_slice(batch);
1338    shape
1339}
1340
1341fn vector_shape(len: SymDim, batch: &[SymDim]) -> Vec<SymDim> {
1342    let mut shape = vec![len];
1343    shape.extend_from_slice(batch);
1344    shape
1345}
1346
1347fn eig_output_dtype(dtype: DType) -> DType {
1348    match dtype {
1349        DType::F64 | DType::C64 => DType::C64,
1350        DType::F32 | DType::C32 => DType::C32,
1351        DType::I32 | DType::I64 | DType::Bool => DType::C64,
1352    }
1353}
1354
1355fn singular_values_dtype(dtype: DType) -> DType {
1356    match dtype {
1357        DType::C64 => DType::F64,
1358        DType::C32 => DType::F32,
1359        other => other,
1360    }
1361}
1362
1363fn promote_dtypes(dtypes: &[DType]) -> DType {
1364    dtypes
1365        .iter()
1366        .copied()
1367        .reduce(tenferro_tensor::validate::promote_dtype)
1368        .unwrap_or(DType::F64)
1369}
1370
1371fn hash_dtype(hasher: &mut dyn Hasher, dtype: DType) {
1372    let tag = match dtype {
1373        DType::F64 => 0,
1374        DType::F32 => 1,
1375        DType::I64 => 2,
1376        DType::C64 => 3,
1377        DType::C32 => 4,
1378        DType::I32 => 5,
1379        DType::Bool => 6,
1380    };
1381    hasher.write_u8(tag);
1382}
1383
1384fn hash_svd_gauge(hasher: &mut dyn Hasher, gauge: SvdGauge) {
1385    let tag = match gauge {
1386        SvdGauge::Raw => 0,
1387        SvdGauge::CanonicalPivot => 1,
1388    };
1389    hasher.write_u8(tag);
1390}
1391
1392fn hash_eigh_gauge(hasher: &mut dyn Hasher, gauge: EighGauge) {
1393    let tag = match gauge {
1394        EighGauge::Raw => 0,
1395        EighGauge::CanonicalPivot => 1,
1396    };
1397    hasher.write_u8(tag);
1398}
1399
1400fn hash_qr_gauge(hasher: &mut dyn Hasher, gauge: QrGauge) {
1401    let tag = match gauge {
1402        QrGauge::Raw => 0,
1403        QrGauge::PositiveDiagonal => 1,
1404    };
1405    hasher.write_u8(tag);
1406}
tenferro_linalg/extension.rs

tenferro_linalg/
extension.rs
