Add dropout test (#265)

* Add dropout test

* Remove condition where training is false

* Skip dropout test when on GPU

* Revert changes to tensor.py and fix test case

* Revert change on whitespace

* Convert Tensor to cpu for testing

* Fix whitespace in tensor.py

test/test_tensor.py

@@ -64,6 +64,14 @@ class TestTinygrad(unittest.TestCase):
     for x,y in zip(test_tinygrad(), test_pytorch()):
       np.testing.assert_allclose(x, y, atol=1e-5)
 
+  def test_dropout(self):
+    Tensor.training = True
+    n, rate = 1_000_000, 0.1
+    w = Tensor.ones(n).dropout(rate)
+    non_zeros = np.count_nonzero(w.cpu().data)
+    expected = n * (1 - rate)
+    np.testing.assert_allclose(non_zeros, expected, rtol=1e-3)
+
   @unittest.skipUnless(not DEFAULT_DEVICE, "float64 not supported on GPU")
   def test_jacobian(self):
     W = np.random.RandomState(1337).random((10, 5))

tinygrad/tensor.py

@@ -267,7 +267,6 @@ class Tensor:
     return self - ss
 
   def dropout(self, p=0.5):
-    # TODO: this needs a test
     if Tensor.training:
       _mask = np.asarray(np.random.binomial(1, 1.0-p, size=self.shape), dtype=self.dtype)
       return self * Tensor(_mask, requires_grad=False, device=self.device) * (1/(1.0 - p))