use candle_core::Tensor;
use nalgebra::{DMatrix, DVector};
pub trait EssParam: Clone {
fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self;
}
pub trait EssParamSummary: EssParam {
fn as_slice(&self) -> &[f32];
fn dim(&self) -> usize;
}
impl EssParam for DVector<f32> {
fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self {
self * a + other * b
}
}
impl EssParamSummary for DVector<f32> {
fn as_slice(&self) -> &[f32] {
self.as_slice()
}
fn dim(&self) -> usize {
self.nrows()
}
}
impl EssParam for DMatrix<f32> {
fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self {
self * a + other * b
}
}
impl EssParamSummary for DMatrix<f32> {
fn as_slice(&self) -> &[f32] {
self.as_slice()
}
fn dim(&self) -> usize {
self.nrows() * self.ncols()
}
}
impl<P: EssParam> EssParam for Vec<P> {
fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self {
self.iter()
.zip(other.iter())
.map(|(s, o)| s.linear_combine(a, o, b))
.collect()
}
}
impl EssParam for Tensor {
fn linear_combine(&self, a: f32, other: &Self, b: f32) -> Self {
(self * a as f64)
.unwrap()
.add(&(other * b as f64).unwrap())
.unwrap()
}
}