#![allow(dead_code, unused_imports)]
use std::sync::Arc;
use std::sync::{Mutex, OnceLock};
use std::{ffi::OsString, sync::MutexGuard};
use num_complex::{Complex32, Complex64};
#[cfg(feature = "cpu-faer")]
use super::linalg::faer_linalg;
use crate::LinalgBackend;
#[cfg(feature = "cpu-faer")]
use tenferro_cpu::linalg_interop::BufferPool;
use tenferro_cpu::{
abs, add, broadcast_in_dim, clamp, compare, conj, div, dynamic_slice, dynamic_update_slice,
embed_diagonal, extract_diagonal, gather, maximum, minimum, mul, neg, pad, reduce_max,
reduce_min, reduce_prod, reduce_sum, reshape, scatter, select, sign, transpose, tril, triu,
CpuBackend, CpuContext,
};
use tenferro_tensor::backend::TensorBackend;
use tenferro_tensor::config::{
CompareDir, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
};
use tenferro_tensor::types::{DType, Tensor, TensorRead, TensorView, TypedTensor};
fn get_f64(t: &Tensor, idx: &[usize]) -> f64 {
match t {
Tensor::F64(inner) => *inner.get(idx).unwrap(),
_ => panic!("expected F64 tensor"),
}
}
fn get_c64(t: &Tensor, idx: &[usize]) -> Complex64 {
match t {
Tensor::C64(inner) => *inner.get(idx).unwrap(),
_ => panic!("expected C64 tensor"),
}
}
fn get_f32(t: &Tensor, idx: &[usize]) -> f32 {
match t {
Tensor::F32(inner) => *inner.get(idx).unwrap(),
_ => panic!("expected F32 tensor"),
}
}
fn get_c32(t: &Tensor, idx: &[usize]) -> Complex32 {
match t {
Tensor::C32(inner) => *inner.get(idx).unwrap(),
_ => panic!("expected C32 tensor"),
}
}
fn get_i64(t: &Tensor, idx: &[usize]) -> i64 {
match t {
Tensor::I64(inner) => *inner.get(idx).unwrap(),
_ => panic!("expected I64 tensor"),
}
}
fn get_i32(t: &Tensor, idx: &[usize]) -> i32 {
match t {
Tensor::I32(inner) => *inner.get(idx).unwrap(),
_ => panic!("expected I32 tensor"),
}
}
fn get_bool(t: &Tensor, idx: &[usize]) -> bool {
match t {
Tensor::Bool(inner) => *inner.get(idx).unwrap(),
_ => panic!("expected Bool tensor"),
}
}
fn assert_f64_close(actual: f64, expected: f64) {
assert!(
(actual - expected).abs() < 1.0e-12,
"expected {expected}, got {actual}"
);
}
fn assert_f64_close_tol(actual: f64, expected: f64, tol: f64) {
assert!(
(actual - expected).abs() < tol,
"expected {expected}, got {actual}, tol={tol}"
);
}
fn assert_c64_close(actual: Complex64, expected: Complex64) {
assert_f64_close(actual.re, expected.re);
assert_f64_close(actual.im, expected.im);
}
fn assert_c64_close_tol(actual: Complex64, expected: Complex64, tol: f64) {
assert_f64_close_tol(actual.re, expected.re, tol);
assert_f64_close_tol(actual.im, expected.im, tol);
}
fn col_major_index(rows: usize, row: usize, col: usize) -> usize {
row + col * rows
}
fn transpose_f64(mat: &[f64], rows: usize, cols: usize) -> Vec<f64> {
let mut out = vec![0.0; rows * cols];
for j in 0..cols {
for i in 0..rows {
out[col_major_index(cols, j, i)] = mat[col_major_index(rows, i, j)];
}
}
out
}
fn transpose_c64(mat: &[Complex64], rows: usize, cols: usize) -> Vec<Complex64> {
let mut out = vec![Complex64::new(0.0, 0.0); rows * cols];
for j in 0..cols {
for i in 0..rows {
out[col_major_index(cols, j, i)] = mat[col_major_index(rows, i, j)];
}
}
out
}
fn conjugate_transpose_c64(mat: &[Complex64], rows: usize, cols: usize) -> Vec<Complex64> {
let mut out = vec![Complex64::new(0.0, 0.0); rows * cols];
for j in 0..cols {
for i in 0..rows {
out[col_major_index(cols, j, i)] = mat[col_major_index(rows, i, j)].conj();
}
}
out
}
fn matmul_f64(lhs: &[f64], rhs: &[f64], m: usize, k: usize, n: usize) -> Vec<f64> {
let mut out = vec![0.0; m * n];
for j in 0..n {
for p in 0..k {
let rhs_pj = rhs[col_major_index(k, p, j)];
for i in 0..m {
out[col_major_index(m, i, j)] += lhs[col_major_index(m, i, p)] * rhs_pj;
}
}
}
out
}
fn matmul_c64(
lhs: &[Complex64],
rhs: &[Complex64],
m: usize,
k: usize,
n: usize,
) -> Vec<Complex64> {
let mut out = vec![Complex64::new(0.0, 0.0); m * n];
for j in 0..n {
for p in 0..k {
let rhs_pj = rhs[col_major_index(k, p, j)];
for i in 0..m {
out[col_major_index(m, i, j)] += lhs[col_major_index(m, i, p)] * rhs_pj;
}
}
}
out
}
fn diag_f64(values: &[f64]) -> Vec<f64> {
let mut out = vec![0.0; values.len() * values.len()];
for (i, value) in values.iter().enumerate() {
out[col_major_index(values.len(), i, i)] = *value;
}
out
}
fn diag_c64(values: &[Complex64]) -> Vec<Complex64> {
let mut out = vec![Complex64::new(0.0, 0.0); values.len() * values.len()];
for (i, value) in values.iter().enumerate() {
out[col_major_index(values.len(), i, i)] = *value;
}
out
}
fn diag_c64_from_real(values: &[f64]) -> Vec<Complex64> {
let mut out = vec![Complex64::new(0.0, 0.0); values.len() * values.len()];
for (i, value) in values.iter().enumerate() {
out[col_major_index(values.len(), i, i)] = Complex64::new(*value, 0.0);
}
out
}
fn batch_matrix_f64_from_tensor(
t: &Tensor,
rows: usize,
cols: usize,
batch_idx: usize,
) -> Vec<f64> {
let mut out = vec![0.0; rows * cols];
for j in 0..cols {
for i in 0..rows {
out[col_major_index(rows, i, j)] = get_f64(t, &[i, j, batch_idx]);
}
}
out
}
fn matrix_f64_from_tensor(t: &Tensor, rows: usize, cols: usize) -> Vec<f64> {
let mut out = vec![0.0; rows * cols];
for j in 0..cols {
for i in 0..rows {
out[col_major_index(rows, i, j)] = get_f64(t, &[i, j]);
}
}
out
}
fn batch_vector_f64_from_tensor(t: &Tensor, len: usize, batch_idx: usize) -> Vec<f64> {
let mut out = vec![0.0; len];
for (i, value) in out.iter_mut().enumerate().take(len) {
*value = get_f64(t, &[i, batch_idx]);
}
out
}
fn vector_f64_from_tensor(t: &Tensor, len: usize) -> Vec<f64> {
let mut out = vec![0.0; len];
for (i, value) in out.iter_mut().enumerate().take(len) {
*value = get_f64(t, &[i]);
}
out
}
fn matrix_c64_from_tensor(t: &Tensor, rows: usize, cols: usize) -> Vec<Complex64> {
let mut out = vec![Complex64::new(0.0, 0.0); rows * cols];
for j in 0..cols {
for i in 0..rows {
out[col_major_index(rows, i, j)] = get_c64(t, &[i, j]);
}
}
out
}
fn batch_matrix_c64_from_tensor(
t: &Tensor,
rows: usize,
cols: usize,
batch_idx: usize,
) -> Vec<Complex64> {
let mut out = vec![Complex64::new(0.0, 0.0); rows * cols];
for j in 0..cols {
for i in 0..rows {
out[col_major_index(rows, i, j)] = get_c64(t, &[i, j, batch_idx]);
}
}
out
}
fn vector_c64_from_tensor(t: &Tensor, len: usize) -> Vec<Complex64> {
let mut out = vec![Complex64::new(0.0, 0.0); len];
for (i, value) in out.iter_mut().enumerate().take(len) {
*value = get_c64(t, &[i]);
}
out
}
fn batch_vector_c64_from_tensor(t: &Tensor, len: usize, batch_idx: usize) -> Vec<Complex64> {
let mut out = vec![Complex64::new(0.0, 0.0); len];
for (i, value) in out.iter_mut().enumerate().take(len) {
*value = get_c64(t, &[i, batch_idx]);
}
out
}
fn simple_gather_config() -> GatherConfig {
GatherConfig {
offset_dims: vec![],
collapsed_slice_dims: vec![0],
start_index_map: vec![0],
index_vector_dim: 1,
slice_sizes: vec![1],
}
}
fn diagonal_scatter_config() -> ScatterConfig {
ScatterConfig {
update_window_dims: vec![],
inserted_window_dims: vec![0, 1],
scatter_dims_to_operand_dims: vec![0, 1],
index_vector_dim: 1,
}
}
mod backend;
mod dtype;
mod linalg;