use flashlight_tensor::tensor::Tensor;
use rand::{rng, Rng};
pub struct Dropout{
dropout: f32,
}
impl Dropout{
pub fn new(dropout: f32) -> Option<Self>{
if dropout>1.0 || dropout < 0.0{
return None;
}
Some(Self{dropout})
}
pub fn forward(&self, tensor: Tensor<f32>) -> Tensor<f32>{
let mut tensor_copy = tensor.clone();
let mut rng = rand::rng();
let neuron_size = tensor_copy.get_shape()[1];
let dropped_neurons = ((neuron_size as f32) * self.dropout) as u32;
for sample in 0..tensor_copy.get_shape()[0]{
for i in 0..dropped_neurons{
let rand_neuron = rng.random_range(0..neuron_size);
tensor_copy.set(0.0, &[sample, rand_neuron]);
}
}
tensor_copy
}
}