From 95f5c85bf3f5fcdb11dd3f1e4d41312257bc8301 Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Mon, 20 Jul 2026 21:43:23 -0400
Subject: [PATCH] some realize and corealize for slow tests (#17106)

---
 test/unit/test_hashing.py    | 2 +-
 test/unit/test_linalg.py     | 7 ++++++-
 test/unit/test_randomness.py | 2 +-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/test/unit/test_hashing.py b/test/unit/test_hashing.py
index 62b75093a0..a3924e655e 100644
--- a/test/unit/test_hashing.py
+++ b/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"))
diff --git a/test/unit/test_linalg.py b/test/unit/test_linalg.py
index 0fd3a8f5bd..1b04303c1f 100644
--- a/test/unit/test_linalg.py
+++ b/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()
diff --git a/test/unit/test_randomness.py b/test/unit/test_randomness.py
index 0d6aae0ca8..0383059749 100644
--- a/test/unit/test_randomness.py
+++ b/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: