use crate::error::Error;
use crate::neural_network::Tensor;
use crate::neural_network::losses::validate_same_shape;
use crate::neural_network::traits::Loss;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct MeanSquaredError;
impl MeanSquaredError {
pub fn new() -> Self {
Self {}
}
}
impl Loss for MeanSquaredError {
fn compute_loss(&self, y_true: &Tensor, y_pred: &Tensor) -> Result<f32, Error> {
validate_same_shape(y_true, y_pred)?;
let squared_diff = (y_pred - y_true).mapv(|x| x * x);
let n = squared_diff.len() as f32;
Ok(squared_diff.sum() / n)
}
fn compute_grad(&self, y_true: &Tensor, y_pred: &Tensor) -> Result<Tensor, Error> {
validate_same_shape(y_true, y_pred)?;
let diff = y_pred - y_true;
let n = diff.len() as f32;
let mut result = diff.clone();
result.par_mapv_inplace(|x| 2.0 * x / n);
Ok(result)
}
}