mawei/tinygrad
1
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-13 00:15:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

add test_embedding to test_softmax_fusion (#9832)

Browse Source
This commit is contained in:
chenyu
2025-04-10 08:25:34 -04:00
committed by GitHub
parent 995d20673a
commit 7fa5f29582
1 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
test/test_softmax_fusion.py
View File

@@ -61,6 +61,20 @@ class TestFuse(unittest.TestCase):
c = (Tensor.rand(N,N)-0.5).realize()
self._test_fuse(lambda a,b,c: a@b@c, a, b, c, atol=1e-5)
def test_embedding(self):
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
vocab_sz = 123
embed_sz = 16
weight = (Tensor.rand(vocab_sz, embed_sz)-0.5).realize()
a = Tensor([1, 1, 2, 3]).realize()
def embedding(idx:Tensor):
arange = Tensor.arange(vocab_sz).unsqueeze(-1)
big_shp = idx.shape + (vocab_sz, embed_sz)
arange, vals = arange.expand(big_shp), weight.expand(big_shp)
idx = idx.reshape(idx.shape+(1, 1)).expand(big_shp)
return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
self._test_fuse(embedding, a, atol=1e-5)
@unittest.skip("still broken")
def test_flash_attention(self):
BS = 4