mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

triple gemm

Browse source
This commit is contained in:
George Hotz 2025-07-30 17:03:02 -07:00
commit 8dff2c1375
1 changed files with 8 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

8
test/test_tiny.py
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):