From b886d250fb97557e3b640ffd3aa25b7576097006 Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Tue, 11 Jun 2024 11:36:02 -0400
Subject: [PATCH] improve test_dropout_on_shard (#4912)

tested some basic property, also minor formatting for a few Tensor.training setups
---
 test/external/external_test_optim.py |  3 ++-
 test/test_multitensor.py             | 11 ++++++-----
 test/test_optim.py                   |  3 ++-
 test/testextra/test_lr_scheduler.py  |  3 ++-
 4 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/test/external/external_test_optim.py b/test/external/external_test_optim.py
index 928c51a1..42abd884 100644
--- a/test/external/external_test_optim.py
+++ b/test/external/external_test_optim.py
@@ -104,7 +104,8 @@ class ExternalTestOptim(unittest.TestCase):
   def setUp(self):
     self.old_training = Tensor.training
     Tensor.training = True
-  def tearDown(self): Tensor.training = self.old_training
+  def tearDown(self):
+    Tensor.training = self.old_training
 
   def _test_optim(self, tinygrad_optim, tensorflow_optim, steps, opts, atol, rtol, tiny_sched=None, tf_sched=None, schedopts=None, do_optim=True):
     for x,y in zip(step(tinygrad_optim, steps=steps, kwargs=opts, scheduler=tiny_sched, schedopts=schedopts, do_optim=do_optim),
diff --git a/test/test_multitensor.py b/test/test_multitensor.py
index 638de7dd..d544c16e 100644
--- a/test/test_multitensor.py
+++ b/test/test_multitensor.py
@@ -504,11 +504,12 @@ class TestMultiTensor(unittest.TestCase):
       t_none.assign(t_zero)
 
   def test_dropout_on_shard(self):
-    Tensor.training = True
-    X = Tensor.ones(256).to(devices_2)
-    output = X.dropout(0.5)
-    output.numpy()
-    Tensor.training = False
+    with Tensor.train():
+      X = Tensor.ones(256).to(devices_2)
+      output = X.dropout(0.5).numpy()
+      unique, counts = np.unique(output, return_counts=True)
+      assert set(unique) == {0, 2}, unique
+      assert 100 < counts[0] < 156, counts[0]
 
   def test_broadcast_const(self):
     devices = (d0, d1, d2, d3)
diff --git a/test/test_optim.py b/test/test_optim.py
index 8501d20d..6e1c3b9f 100644
--- a/test/test_optim.py
+++ b/test/test_optim.py
@@ -46,7 +46,8 @@ class TestOptim(unittest.TestCase):
   def setUp(self):
     self.old_training = Tensor.training
     Tensor.training = True
-  def tearDown(self): Tensor.training = self.old_training
+  def tearDown(self):
+    Tensor.training = self.old_training
 
   def _test_optim(self, tinygrad_optim, torch_optim, steps, opts, atol, rtol):
     for x,y in zip(step(Tensor, tinygrad_optim, steps, **opts),
diff --git a/test/testextra/test_lr_scheduler.py b/test/testextra/test_lr_scheduler.py
index ac921f92..2f125470 100644
--- a/test/testextra/test_lr_scheduler.py
+++ b/test/testextra/test_lr_scheduler.py
@@ -56,7 +56,8 @@ class TestLrScheduler(unittest.TestCase):
   def setUp(self):
     self.old_training = Tensor.training
     Tensor.training = True
-  def tearDown(self): Tensor.training = self.old_training
+  def tearDown(self):
+    Tensor.training = self.old_training
 
   def _test_lr_scheduler(self, tinygrad_sched, torch_sched, epochs, opts, atol, rtol, adam=True):
     accs = opts.pop('accs', None)