tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 17:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

update embedding test to not use dtypes.long [pr] (#9556)

Browse Source
This commit is contained in:
chenyu
2025-03-23 21:33:38 -04:00
committed by GitHub
parent c965f4c20b
commit ba41076e94
2 changed files with 5 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_multitensor.py
+1 -2
View File
@@ -307,12 +307,11 @@ class TestMultiTensor(unittest.TestCase):
lr_sched = OneCycleLR(optim, max_lr=0.1, pct_start=0.1, div_factor=100, final_div_factor=0.1, total_steps=10)
lr_sched.step()
@unittest.skipUnless(is_dtype_supported(dtypes.long), f"long dtype not supported on {Device.DEFAULT}")
def test_embedding(self):
B, T, embed_size, vocab_size = 4, 10, 20, 28
layer = nn.Embedding(vocab_size, embed_size)
x = Tensor(np.random.randint(0, vocab_size, (B, T)))
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
z = layer(x)
layer_sharded = nn.Embedding(vocab_size, embed_size)
test/test_nn.py
+4 -6
View File
@@ -2,14 +2,13 @@
import unittest
import numpy as np
import torch
from tinygrad import Tensor, Device, TinyJit, dtypes
from tinygrad import Tensor, Device, TinyJit
from tinygrad.ops import Ops
from tinygrad.helpers import CI, Context, OSX
from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
from tinygrad.nn.state import load_state_dict
from tinygrad.engine.realize import run_schedule
from tinygrad.device import is_dtype_supported
@unittest.skipIf(CI and Device.DEFAULT in {"CUDA", "NV"}, "slow")
class TestNN(unittest.TestCase):
@@ -477,7 +476,6 @@ class TestNN(unittest.TestCase):
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=2e-3, rtol=1e-3)
@unittest.skipUnless(is_dtype_supported(dtypes.long), f"no long on {Device.DEFAULT}")
def test_embedding(self):
B, T, embed_size, vocab_size = 4, 10, 20, 28
@@ -489,14 +487,14 @@ class TestNN(unittest.TestCase):
torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
# test
x = Tensor(np.random.randint(0, vocab_size, (B, T)))
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
z = layer(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
# test with empty input length
x = Tensor(np.random.randint(0, vocab_size, (B, 0)))
x = Tensor(np.random.randint(0, vocab_size, (B, 0), dtype=np.int32))
z = layer(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
@@ -508,7 +506,7 @@ class TestNN(unittest.TestCase):
return layer(x).realize()
for _ in range(3):
x = Tensor(np.random.randint(0, vocab_size, (B, T)))
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
z = layer_jit(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)