some realize and corealize for slow tests (#17106)

test/unit/test_hashing.py

@@ -71,7 +71,7 @@ class TestKeccak(unittest.TestCase):
   def test_variable_bs(self):
     data = Tensor([b"abc", b"abc", b"def"], dtype=dtypes.uint8).repeat(2048, 1)
     bs = UOp.variable("bs", 1, 4096).bind(3)
-    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(3, 32)
+    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(3, 32).realize()
     self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
     self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
     self.assertEqual(bytes(out[2].tolist()), bytearray.fromhex("8e0d8f672252acb0 ffc5093db8653b18 1513bf9a2097e737 b4f73533dcaf46df"))

test/unit/test_linalg.py

@@ -17,6 +17,7 @@ class TestLinAlg(unittest.TestCase):
     for size in sizes:
       a = Tensor.randn(size).realize()
       U,S,V = a.svd()
+      Tensor.realize(U,S,V)
       b_shape,m,n = size[0:-2],size[-2],size[-1]
       k = min(m,n)
       s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
@@ -29,6 +30,7 @@ class TestLinAlg(unittest.TestCase):
     with Context(CHECK_OOB=0):  # sometimes this is slow in CI
       a = Tensor.randn(size).realize()
       U,S,V = a.svd(full_matrices=False)
+      Tensor.realize(U,S,V)
       b_shape,m,n = size[0:-2],size[-2],size[-1]
       k = min(m,n)
       s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
@@ -61,6 +63,7 @@ class TestLinAlg(unittest.TestCase):
     for size in sizes:
       a = Tensor.randn(size).realize()
       Q,R = a.qr()
+      Tensor.realize(Q,R)
       orthogonality_helper(Q)
       reconstruction_helper([Q,R],a)
 
@@ -73,9 +76,10 @@ class TestLinAlg(unittest.TestCase):
     reconstruction_helper([Q,R], a)
 
   def test_svd_identity(self):
-    for a in (Tensor.eye(2), Tensor.zeros(2, 2)):
+    for a in (Tensor.eye(2).clone(), Tensor.zeros(2, 2)):
       a = a.realize()
       U,S,V = a.svd()
+      Tensor.realize(U,S,V)
       assert not np.isnan(U.numpy()).any()
       assert not np.isnan(S.numpy()).any()
       assert not np.isnan(V.numpy()).any()
@@ -85,6 +89,7 @@ class TestLinAlg(unittest.TestCase):
   def test_svd_identity_4x4(self):
     a = Tensor.eye(4).clone()
     U,S,V = a.svd()
+    Tensor.realize(U,S,V)
     assert not np.isnan(U.numpy()).any()
     assert not np.isnan(S.numpy()).any()
     assert not np.isnan(V.numpy()).any()

test/unit/test_randomness.py

@@ -119,7 +119,7 @@ class TestRandomness(unittest.TestCase):
     self.assertRaises(AssertionError, lambda: Tensor(2).multinomial(1, replacement=False))
     self.assertRaises(AssertionError, lambda: Tensor([1, 9]).multinomial(0, replacement=False))
     def _check_with_torch(w, num_samples, replacement):
-      tiny_res = Tensor(w).multinomial(num_samples, replacement=replacement)
+      tiny_res = Tensor(w).multinomial(num_samples, replacement=replacement).realize()
       torch_res = torch.tensor(w).multinomial(num_samples, replacement=replacement)
       self.assertEqual(tiny_res.shape, torch_res.shape)
       if torch_res.ndim == 1: