rotta_rs 0.0.5

a Deep Learning library with rust language
Documentation
use std::sync::{ Arc, Mutex };

use rand::rngs::StdRng;
use rand_distr::{ Bernoulli, Distribution };

use crate::rotta_rs_module::{ arrayy::Arrayy, Tensor };

pub struct Dropout {
    pub bernoulli: Bernoulli,
    pub p: f64,
    pub rng: StdRng,
    pub eval: Arc<Mutex<bool>>,
}

impl Dropout {
    pub fn forward(&mut self, x: &Tensor) -> Tensor {
        if !self.eval_status() {
            let r = self.r(x.shape());
            r.set_requires_grad(false);

            x * &r
        } else {
            (*x).clone()
        }
    }

    fn r(&mut self, shape: Vec<usize>) -> Tensor {
        let arr = Arrayy::arrayy_from_shape_fn(shape, || {
            let prob = self.bernoulli.sample(&mut self.rng);
            if prob {
                0.0
            } else {
                1.0 / (1.0 - self.p)
            }
        });
        Tensor::from_arrayy(arr)
    }

    pub fn eval_status(&self) -> bool {
        *self.eval.lock().unwrap()
    }

    pub fn eval(&self) {
        *self.eval.lock().unwrap() = true;
    }

    pub fn train(&self) {
        *self.eval.lock().unwrap() = false;
    }
}