use crate::rotta_rs_module::{ arrayy::Arrayy, BackwardLabel, NodeType, Tensor };
#[allow(dead_code)]
pub fn relu(x: &Tensor) -> Tensor {
let value = x.value();
let output = value.map(|x| {
if *x >= 0.0 { *x } else { 0.0 }
});
let tensor = Tensor::from_arrayy(output);
tensor.update_parent(vec![x.node.clone()]);
tensor.node.lock().as_mut().unwrap().label = Some(BackwardLabel::Relu(x.node.clone()));
tensor
}
#[allow(dead_code)]
pub fn d_relu(x: &NodeType, grad: &Arrayy) {
let mut x_lock = x.lock().unwrap();
if x_lock.requires_grad {
let d_x =
x_lock.value.map(|x| {
if *x >= 0.0 { 1.0 } else { 0.0 }
}) * grad;
x_lock.add_grad(d_x);
}
}