From 8dff2c1375fa559c7c7c0fa8048e4bcdd4cc452f Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Wed, 30 Jul 2025 17:03:02 -0700
Subject: [PATCH] triple gemm

---
 test/test_tiny.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/test/test_tiny.py b/test/test_tiny.py
index a38c7ed628..265f25656e 100644
--- a/test/test_tiny.py
+++ b/test/test_tiny.py
@@ -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):