sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

triple gemm

Browse Source
This commit is contained in:
geohot
2025-07-30 17:03:02 -07:00
parent 3cff1a6b13
commit 8dff2c1375
1 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_tiny.py
+8
View File
@@ -33,6 +33,14 @@ class TestTiny(unittest.TestCase):
self.assertListEqual((out:=a@b).flatten().tolist(), [1.0]*(N*N))
if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
def test_double_gemm(self, N=64, out_dtype=dtypes.float):
a = Tensor.ones(N,N).contiguous().realize()
b = Tensor.eye(N).contiguous().realize()
c = Tensor.eye(N).contiguous().realize()
d = Tensor.eye(N).contiguous().realize()
out = (((a@b).relu()@c).relu()@d).contiguous().realize()
self.assertListEqual(out.flatten().tolist(), [1.0]*(N*N))
# *** randomness ***
def test_random(self):