mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-06-11 23:46:02 +08:00
41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
import unittest
|
|
from tinygrad import Tensor, Context, Variable, Device
|
|
from test.helpers import needs_second_gpu
|
|
|
|
class TestValidateWithCPU(unittest.TestCase):
|
|
def setUp(self):
|
|
self.ctx = Context(VALIDATE_WITH_CPU=1)
|
|
self.ctx.__enter__()
|
|
def tearDown(self): self.ctx.__exit__(None, None, None)
|
|
|
|
def test_add(self): self.assertListEqual((Tensor([1.,2,3])+Tensor([4.,5,6])).tolist(), [5.0, 7.0, 9.0])
|
|
def test_mul(self): self.assertListEqual((Tensor([1.,2,3])*Tensor([4.,5,6])).tolist(), [4.0, 10.0, 18.0])
|
|
def test_sum(self): self.assertEqual(Tensor([1.,2,3,4]).sum().item(), 10.0)
|
|
def test_reduce_then_op(self): self.assertEqual((Tensor([1.,2,3,4]).sum() * 2).item(), 20.0)
|
|
|
|
def test_assign(self):
|
|
a = Tensor([1.,2,3]).realize()
|
|
a.assign(a + 1).realize()
|
|
self.assertListEqual(a.tolist(), [2.0, 3.0, 4.0])
|
|
|
|
def test_buffer_view(self):
|
|
self.assertListEqual((Tensor([1.,2,3,4,5,6,7,8])[2:6] + 1).tolist(), [4.0, 5.0, 6.0, 7.0])
|
|
|
|
def test_symbolic(self):
|
|
i = Variable('i', 1, 10)
|
|
ones = Tensor.ones(10).contiguous()
|
|
self.assertListEqual((ones[:i.bind(5)] + 1).contiguous()[:5].tolist(), [2.0]*5)
|
|
|
|
def test_multi_kernel(self):
|
|
a = (Tensor([1.,2,3]) + 1).contiguous()
|
|
b = (a * 2).contiguous()
|
|
self.assertListEqual((b - 1).tolist(), [3.0, 5.0, 7.0])
|
|
|
|
@needs_second_gpu
|
|
def test_sharded(self):
|
|
t = Tensor([1.,2,3,4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0)
|
|
self.assertListEqual((t + 1).tolist(), [2.0, 3.0, 4.0, 5.0])
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|