use std::sync::{Arc, OnceLock};
use tenferro_ad::error::{Error, Result};
use tenferro_ad::extension::apply_eager_with_extension_session;
use tenferro_ad::EagerTensor;
use tenferro_runtime::{ErrorPhase, ExtensionModule};
use crate::eager_composites;
use crate::extension::{
execute_linalg_extension_reads, extension_module, validate_derivative_eps, EighOptions,
LinalgExtensionOp, LinalgOp, QrOptions, SvdOptions,
};
pub trait EagerTensorLinalgExt {
fn svd(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)>;
fn svd_with_options(
&self,
options: SvdOptions,
) -> Result<(EagerTensor, EagerTensor, EagerTensor)>;
fn svd_full(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)>;
fn qr(&self) -> Result<(EagerTensor, EagerTensor)>;
fn qr_with_options(&self, options: QrOptions) -> Result<(EagerTensor, EagerTensor)>;
fn lu(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor, EagerTensor)>;
fn full_piv_lu(
&self,
) -> Result<(
EagerTensor,
EagerTensor,
EagerTensor,
EagerTensor,
EagerTensor,
)>;
fn full_piv_lu_solve(&self, b: &EagerTensor) -> Result<EagerTensor>;
fn solve(&self, b: &EagerTensor) -> Result<EagerTensor>;
fn lstsq(&self, b: &EagerTensor) -> Result<EagerTensor>;
fn cholesky(&self) -> Result<EagerTensor>;
fn eigh(&self) -> Result<(EagerTensor, EagerTensor)>;
fn eigh_with_options(&self, options: EighOptions) -> Result<(EagerTensor, EagerTensor)>;
fn eig(&self) -> Result<(EagerTensor, EagerTensor)>;
fn triangular_solve(
&self,
b: &EagerTensor,
left_side: bool,
lower: bool,
transpose_a: bool,
unit_diagonal: bool,
) -> Result<EagerTensor>;
fn slogdet(&self) -> Result<(EagerTensor, EagerTensor)>;
fn det(&self) -> Result<EagerTensor>;
fn inv(&self) -> Result<EagerTensor>;
fn eigvalsh(&self) -> Result<EagerTensor>;
fn eigvals(&self) -> Result<EagerTensor>;
fn pinv(&self) -> Result<EagerTensor>;
fn pinv_with_rtol(&self, rtol: f64) -> Result<EagerTensor>;
fn norm(&self, ord: Option<f64>, dim: Option<&[usize]>, keepdim: bool) -> Result<EagerTensor>;
}
impl EagerTensorLinalgExt for EagerTensor {
fn svd(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
svd(self)
}
fn svd_with_options(
&self,
options: SvdOptions,
) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
svd_with_options(self, options)
}
fn svd_full(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
svd_full(self)
}
fn qr(&self) -> Result<(EagerTensor, EagerTensor)> {
qr(self)
}
fn qr_with_options(&self, options: QrOptions) -> Result<(EagerTensor, EagerTensor)> {
qr_with_options(self, options)
}
fn lu(&self) -> Result<(EagerTensor, EagerTensor, EagerTensor, EagerTensor)> {
lu(self)
}
fn full_piv_lu(
&self,
) -> Result<(
EagerTensor,
EagerTensor,
EagerTensor,
EagerTensor,
EagerTensor,
)> {
full_piv_lu(self)
}
fn full_piv_lu_solve(&self, b: &EagerTensor) -> Result<EagerTensor> {
full_piv_lu_solve(self, b)
}
fn solve(&self, b: &EagerTensor) -> Result<EagerTensor> {
solve(self, b)
}
fn lstsq(&self, b: &EagerTensor) -> Result<EagerTensor> {
eager_composites::lstsq(self, b)
}
fn cholesky(&self) -> Result<EagerTensor> {
cholesky(self)
}
fn eigh(&self) -> Result<(EagerTensor, EagerTensor)> {
eigh(self)
}
fn eigh_with_options(&self, options: EighOptions) -> Result<(EagerTensor, EagerTensor)> {
eigh_with_options(self, options)
}
fn eig(&self) -> Result<(EagerTensor, EagerTensor)> {
eig(self)
}
fn triangular_solve(
&self,
b: &EagerTensor,
left_side: bool,
lower: bool,
transpose_a: bool,
unit_diagonal: bool,
) -> Result<EagerTensor> {
triangular_solve(self, b, left_side, lower, transpose_a, unit_diagonal)
}
fn slogdet(&self) -> Result<(EagerTensor, EagerTensor)> {
eager_composites::slogdet(self)
}
fn det(&self) -> Result<EagerTensor> {
eager_composites::det(self)
}
fn inv(&self) -> Result<EagerTensor> {
eager_composites::inv(self)
}
fn eigvalsh(&self) -> Result<EagerTensor> {
eager_composites::eigvalsh(self)
}
fn eigvals(&self) -> Result<EagerTensor> {
eager_composites::eigvals(self)
}
fn pinv(&self) -> Result<EagerTensor> {
eager_composites::pinv(self)
}
fn pinv_with_rtol(&self, rtol: f64) -> Result<EagerTensor> {
eager_composites::pinv_with_rtol(self, rtol)
}
fn norm(&self, ord: Option<f64>, dim: Option<&[usize]>, keepdim: bool) -> Result<EagerTensor> {
eager_composites::norm(self, ord, dim, keepdim)
}
}
fn apply_linalg_eager(op: LinalgOp, inputs: &[&EagerTensor]) -> Result<Vec<EagerTensor>> {
let op = Arc::new(LinalgExtensionOp::new(op));
let execute_op = Arc::clone(&op);
apply_eager_with_extension_session(
op,
inputs,
eager_cpu_extension_module()?,
move |_op, input_reads, ctx| execute_linalg_extension_reads(&execute_op, input_reads, ctx),
)
}
fn eager_cpu_extension_module() -> Result<Arc<dyn ExtensionModule>> {
static MODULE: OnceLock<Arc<dyn ExtensionModule>> = OnceLock::new();
if let Some(module) = MODULE.get() {
return Ok(Arc::clone(module));
}
let engine_id = tenferro_cpu::runtime_engine_id().map_err(eager_runtime_config_error)?;
let module = extension_module::<tenferro_cpu::CpuBackend>(engine_id)
.map_err(eager_runtime_config_error)?;
let _ = MODULE.set(Arc::clone(&module));
Ok(MODULE.get().cloned().unwrap_or(module))
}
fn eager_runtime_config_error(source: tenferro_runtime::RuntimeConfigError) -> Error {
Error::runtime_state_source(
"tenferro_linalg::eager_extension_module",
ErrorPhase::Execution,
source,
)
}
pub fn svd(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
svd_with_options(a, SvdOptions::default())
}
pub fn svd_with_options(
a: &EagerTensor,
options: SvdOptions,
) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
validate_derivative_eps("svd_with_options", options.derivative_eps)?;
let mut outputs = apply_linalg_eager(
LinalgOp::Svd {
derivative_eps: options.derivative_eps,
gauge: options.gauge,
},
&[a],
)?
.into_iter();
match (
outputs.next(),
outputs.next(),
outputs.next(),
outputs.next(),
) {
(Some(u), Some(s), Some(vt), None) => Ok((u, s, vt)),
_ => Err(Error::Internal(
"svd eager op returned an unexpected number of outputs".to_string(),
)),
}
}
pub fn svd_full(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
let mut outputs = apply_linalg_eager(LinalgOp::SvdFull, &[a])?.into_iter();
match (
outputs.next(),
outputs.next(),
outputs.next(),
outputs.next(),
) {
(Some(u), Some(s), Some(vh), None) => Ok((u, s, vh)),
_ => Err(Error::Internal(
"svd_full eager op returned an unexpected number of outputs".to_string(),
)),
}
}
pub fn qr(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor)> {
qr_with_options(a, QrOptions::default())
}
pub fn qr_with_options(a: &EagerTensor, options: QrOptions) -> Result<(EagerTensor, EagerTensor)> {
two_outputs(
apply_linalg_eager(
LinalgOp::Qr {
gauge: options.gauge,
},
&[a],
)?,
"qr",
)
}
pub fn lu(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor, EagerTensor, EagerTensor)> {
let mut outputs = apply_linalg_eager(LinalgOp::Lu, &[a])?.into_iter();
match (
outputs.next(),
outputs.next(),
outputs.next(),
outputs.next(),
outputs.next(),
) {
(Some(p), Some(l), Some(u), Some(parity), None) => Ok((p, l, u, parity)),
_ => Err(Error::Internal(
"lu eager op returned an unexpected number of outputs".to_string(),
)),
}
}
pub fn full_piv_lu(
a: &EagerTensor,
) -> Result<(
EagerTensor,
EagerTensor,
EagerTensor,
EagerTensor,
EagerTensor,
)> {
let mut outputs = apply_linalg_eager(LinalgOp::FullPivLu, &[a])?.into_iter();
match (
outputs.next(),
outputs.next(),
outputs.next(),
outputs.next(),
outputs.next(),
outputs.next(),
) {
(Some(p), Some(l), Some(u), Some(q), Some(parity), None) => Ok((p, l, u, q, parity)),
_ => Err(Error::Internal(
"full_piv_lu eager op returned an unexpected number of outputs".to_string(),
)),
}
}
pub fn full_piv_lu_solve(a: &EagerTensor, b: &EagerTensor) -> Result<EagerTensor> {
one_output(
apply_linalg_eager(LinalgOp::FullPivLuSolve { transpose_a: false }, &[a, b])?,
"full_piv_lu_solve",
)
}
pub fn solve(a: &EagerTensor, b: &EagerTensor) -> Result<EagerTensor> {
let mut factor_outputs = apply_linalg_eager(LinalgOp::LuFactor, &[a])?.into_iter();
let (packed_lu, pivots) = match (
factor_outputs.next(),
factor_outputs.next(),
factor_outputs.next(),
factor_outputs.next(),
) {
(Some(packed_lu), Some(pivots), Some(_parity), None) => (packed_lu, pivots),
_ => {
return Err(Error::Internal(
"lu_factor eager op returned an unexpected number of outputs".to_string(),
));
}
};
one_output(
apply_linalg_eager(
LinalgOp::LuSolvePrepared {
transpose_a: false,
conjugate_a: false,
},
&[a, &packed_lu, &pivots, b],
)?,
"solve",
)
}
pub fn cholesky(a: &EagerTensor) -> Result<EagerTensor> {
one_output(apply_linalg_eager(LinalgOp::Cholesky, &[a])?, "cholesky")
}
pub fn eigh(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor)> {
eigh_with_options(a, EighOptions::default())
}
pub fn eigh_with_options(
a: &EagerTensor,
options: EighOptions,
) -> Result<(EagerTensor, EagerTensor)> {
validate_derivative_eps("eigh_with_options", options.derivative_eps)?;
two_outputs(
apply_linalg_eager(
LinalgOp::Eigh {
derivative_eps: options.derivative_eps,
gauge: options.gauge,
},
&[a],
)?,
"eigh",
)
}
pub fn eig(a: &EagerTensor) -> Result<(EagerTensor, EagerTensor)> {
two_outputs(
apply_linalg_eager(
LinalgOp::Eig {
input_dtype: a.dtype(),
},
&[a],
)?,
"eig",
)
}
pub fn triangular_solve(
a: &EagerTensor,
b: &EagerTensor,
left_side: bool,
lower: bool,
transpose_a: bool,
unit_diagonal: bool,
) -> Result<EagerTensor> {
one_output(
apply_linalg_eager(
LinalgOp::TriangularSolve {
left_side,
lower,
transpose_a,
unit_diagonal,
},
&[a, b],
)?,
"triangular_solve",
)
}
fn one_output(outputs: Vec<EagerTensor>, name: &str) -> Result<EagerTensor> {
let mut outputs = outputs.into_iter();
match (outputs.next(), outputs.next()) {
(Some(output), None) => Ok(output),
_ => Err(Error::Internal(format!(
"{name} eager op returned an unexpected number of outputs"
))),
}
}
fn two_outputs(outputs: Vec<EagerTensor>, name: &str) -> Result<(EagerTensor, EagerTensor)> {
let mut outputs = outputs.into_iter();
match (outputs.next(), outputs.next(), outputs.next()) {
(Some(lhs), Some(rhs), None) => Ok((lhs, rhs)),
_ => Err(Error::Internal(format!(
"{name} eager op returned an unexpected number of outputs"
))),
}
}