coreylowman · coreylowman · Jul 25, 2022 · Jul 21, 2022 · Jul 21, 2022 · Jul 21, 2022
diff --git a/src/tensor_ops/impl_dropout.rs b/src/tensor_ops/impl_dropout.rs
@@ -1,6 +1,7 @@
+use super::utils::move_tape_and_add_backward_op;
 use crate::prelude::*;
 use rand::Rng;
-use rand_distr::Standard;
+use rand_distr::{Distribution, Standard};
 
 /// Randomly drops out elements from `t` with probability `p`, and multiplies all elements by `1 / (1 - p)`.
 ///
@@ -14,17 +15,17 @@ pub fn dropout<T: Tensor<Dtype = f32>, R: Rng>(t: T, p: f32, rng: &mut R) -> T {
     } else {
         // `t` owns the tape in this branch, so apply dropout randomly.
         let rinvp = (1.0 - p).recip();
-        map_df_uses_fx(
-            t,
-            move |x| {
-                if rng.sample::<f32, Standard>(Standard) < p {
-                    0.0
-                } else {
-                    x * rinvp
-                }
-            },
-            move |fx| if fx > &0.0 { rinvp } else { 0.0 },
-        )
+        let deriv = T::Device::filled(&mut |d| {
+            let val: f32 = Standard.sample(rng);
+            *d = if val < p { 0.0 } else { rinvp };
+        });
+        let mut result = T::NoTape::zeros();
+        T::Device::addmul(result.mut_data(), t.data(), deriv.as_ref());
+
+        move_tape_and_add_backward_op(t, result, move |t, result, grads| {
+            let (t_grad, result_grad) = grads.mut_and_ref(&t, &result);
+            T::Device::addmul(t_grad, deriv.as_ref(), result_grad);
+        })
     }
 }
 
@@ -79,11 +80,11 @@ mod tests {
     }
 
     #[test]
-    fn test_dropout_1d() {
+    fn test_dropout_1d_with_non_positive_values() {
         let mut rng = StdRng::seed_from_u64(3);
-        let t = Tensor1D::new([1.0, 2.0, 3.0, 4.0, 5.0]);
+        let t = Tensor1D::new([0.0, 2.0, -3.0, -4.0, 0.0]);
         let r = t.trace().dropout(0.5, &mut rng);
-        assert_eq!(r.data(), &[2.0, 0.0, 6.0, 0.0, 0.0]);
+        assert_eq!(r.data(), &[0.0, 0.0, -6.0, 0.0, 0.0]);
         let gradients = r.mean().backward();
         assert_eq!(gradients.ref_gradient(&t), &[0.4, 0.0, 0.4, 0.0, 0.0]);
     }