From ba41076e94600fceda2790eda913dffc2dded34d Mon Sep 17 00:00:00 2001 From: chenyu Date: Sun, 23 Mar 2025 21:33:38 -0400 Subject: [PATCH] update embedding test to not use dtypes.long [pr] (#9556) --- test/test_multitensor.py | 3 +-- test/test_nn.py | 10 ++++------ 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/test/test_multitensor.py b/test/test_multitensor.py index 5935774a84..2065269ddf 100644 --- a/test/test_multitensor.py +++ b/test/test_multitensor.py @@ -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) diff --git a/test/test_nn.py b/test/test_nn.py index 987e7e4f17..0307f6737b 100755 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -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)