use std::marker::PhantomData;
use static_assertions::assert_impl_all;
use crate::{Element, Shape, Symbol, Tape, Tensor, Value};
use super::Module;
assert_impl_all!(Dropout<f64>: Send, Sync);
#[derive(Debug, Clone)]
pub struct Dropout<E> {
mask: Symbol,
_marker: PhantomData<E>,
}
impl<E: Element> Dropout<E> {
pub fn new(tape: &Tape<E>, shape: impl Into<Shape>) -> Self {
let mask = tape.input(Tensor::counted(shape.into(), 1));
Self {
mask: mask.symbol(),
_marker: PhantomData,
}
}
pub fn mask(&self) -> Symbol {
self.mask
}
}
impl<E: Element> Module<E> for Dropout<E> {
fn express<'tape>(&self, input: Value<'tape, E>) -> Value<'tape, E> {
let tape = input.tape();
input * tape.resolve(self.mask)
}
}
#[cfg(test)]
#[path = "tests/dropout_tests.rs"]
mod tests;