Files
tinygrad/test/unit/test_tensor_io.py
leopf f0401e14e8 tar_extract with Tensors (#7853)
* initial

* USTAR, PAX and GNU support + testing

* from_bytes byteorder

* use TarInfo.frombuf

* tensor only usage

* remove contextlib.suppress

* shorter ow,pax

* more tests

* testing length + move tests

* cleanup

* new approach: RawTensorIO

* fix fetch

* enable read test

* cleanup and ignore fix

* fix for python < 3.12

* make it RawIO

* functions

---------

Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com>
Co-authored-by: chenyu <chenyu@fastmail.com>
2024-12-04 17:03:19 +08:00

39 lines
1.2 KiB
Python

import unittest
from tinygrad import Tensor, dtypes
from tinygrad.nn.state import TensorIO
class TestTensorIO(unittest.TestCase):
def test_create(self):
with self.assertRaises(ValueError):
TensorIO(Tensor(b"Hello World").reshape(1, -1))
with self.assertRaises(ValueError):
TensorIO(Tensor([], dtype=dtypes.int64).reshape(1, -1))
def test_seek(self):
t = Tensor(b"Hello World!")
fobj = TensorIO(t)
self.assertEqual(fobj.tell(), 0)
self.assertEqual(fobj.seek(1), 1)
self.assertEqual(fobj.seek(-2, 2), len(t) - 2)
self.assertEqual(fobj.seek(1, 1), len(t) - 1)
self.assertEqual(fobj.seek(10, 1), len(t))
self.assertEqual(fobj.seek(10, 2), len(t))
self.assertEqual(fobj.seek(-10, 0), 0)
def test_read(self):
data = b"Hello World!"
fobj = TensorIO(Tensor(data))
self.assertEqual(fobj.read(1), data[:1])
self.assertEqual(fobj.read(5), data[1:6])
self.assertEqual(fobj.read(100), data[6:])
self.assertEqual(fobj.read(100), b"")
def test_read_nolen(self):
data = b"Hello World!"
fobj = TensorIO(Tensor(data))
fobj.seek(2)
self.assertEqual(fobj.read(), data[2:])
if __name__ == '__main__':
unittest.main()