From cde7a7cda712d22adb1417f95a7914af4fe5c27a Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Wed, 29 May 2024 17:26:24 -0400
Subject: [PATCH] isolate the 134ms kernel in train_gpt2.py (#4773)

133ms on tinybox red with BEAM=2
---
 test/external/external_test_lm_head.py | 10 ++++++++++
 1 file changed, 10 insertions(+)
 create mode 100644 test/external/external_test_lm_head.py

diff --git a/test/external/external_test_lm_head.py b/test/external/external_test_lm_head.py
new file mode 100644
index 0000000000..851e0aa882
--- /dev/null
+++ b/test/external/external_test_lm_head.py
@@ -0,0 +1,10 @@
+from tinygrad import Tensor, nn
+
+if __name__ == "__main__":
+  vocab_size = 50257
+  n_embd = 768
+  lm_head = nn.Linear(n_embd, vocab_size, bias=False)
+  bs = 4
+  seq_len = 1024
+  x = Tensor.rand(bs, seq_len, n_embd)
+  ret = lm_head(x).realize()