mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-07 21:13:41 +08:00
Daniel Koepping
32 lines
1.2 KiB
Python
32 lines
1.2 KiB
Python
#!/usr/bin/env python
|
|
import unittest
|
|
from tinygrad import Tensor
|
|
from tinygrad.nn import Conv2d, Linear, Embedding
|
|
from tinygrad.nn.state import load_state_dict
|
|
|
|
class TestNN(unittest.TestCase):
|
|
def test_conv2d_same_padding_invalid_stride(self):
|
|
self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same')
|
|
|
|
def test_conv2d_same_padding_invalid_padding_str(self):
|
|
self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=1, padding='not_same')
|
|
|
|
def test_embedding_shape(self):
|
|
vocab_size, embed_size = 10, 16
|
|
layer = Embedding(vocab_size, embed_size)
|
|
for rank in range(5):
|
|
shp = (1,) * rank
|
|
a = Tensor([3]).reshape(shp)
|
|
result = layer(a)
|
|
self.assertEqual(result.shape, shp + (embed_size,))
|
|
|
|
def test_load_state_dict_shape_mismatch(self):
|
|
d1, d2 = 2, 4
|
|
layer = Linear(d1, d1, bias=False)
|
|
state_dict = {'weight': Tensor.randn(d2, d2)}
|
|
with self.assertRaisesRegex(ValueError, r'Shape mismatch in layer `weight`: Expected shape \(2, 2\), but found \(4, 4\) in state dict.'):
|
|
load_state_dict(layer, state_dict)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|