Files
sunnypilot-byd/tinygrad_repo/test/external/external_test_lm_head.py
T
github-actions[bot] de77c89cf6 sunnypilot v2026.002.000
version: sunnypilot v2026.002.000 (staging)
date: 2026-07-21T07:44:02
master commit: 2c334ede443d7391d27575af8c854a095ba702a8
2026-07-21 07:44:02 +00:00

11 lines
244 B
Python

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()