forked from tinygrad/tinygrad
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_image_dtype.py
77 lines (66 loc) · 2.78 KB
/
test_image_dtype.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import unittest
import numpy as np
from tinygrad import Device, dtypes, Tensor, Variable
from tinygrad.dtype import ImageDType
from tinygrad.features.image import to_image_idx
@unittest.skipIf(Device.DEFAULT != "GPU", "only images on GPU")
class TestImageDType(unittest.TestCase):
def test_image_and_back(self):
data = Tensor.randn(9*27*4).realize()
tst = data.numpy()
it = data.cast(dtypes.imagef((9,27,4))).realize()
assert isinstance(it.lazydata.base.realized.dtype, ImageDType)
np.testing.assert_equal(tst, it.numpy())
def test_image_and_back_wrong_shape(self):
data = Tensor.randn(9*27*4).realize()
tst = data.numpy()
it = data.cast(dtypes.imagef((9,12,4))).realize()
assert not isinstance(it.lazydata.base.realized.dtype, ImageDType)
np.testing.assert_equal(tst, it.numpy())
def test_shrink_load_float(self):
it = Tensor.randn(4).cast(dtypes.imagef((1,1,4))).realize()
imgv = it.numpy()
np.testing.assert_equal(imgv[0:2], it[0:2].numpy())
def test_mul_stays_image(self):
it = Tensor.randn(4).cast(dtypes.imagef((1,1,4))).realize()
out = (it*2).realize()
assert isinstance(out.lazydata.base.realized.dtype, ImageDType)
def test_shrink_max(self):
it = Tensor.randn(8).cast(dtypes.imagef((1,2,4))).realize()
imgv = it.numpy()
np.testing.assert_equal(np.maximum(imgv[0:3], 0), it[0:3].relu().numpy())
def test_shrink_to_float(self):
it = Tensor.randn(4, 4).cast(dtypes.imagef((1,4,4))).realize()
imgv = it.numpy()
np.testing.assert_equal(np.maximum(imgv[:, 0], 0), it[:, 0].relu().realize())
def test_lru_alloc(self):
data = Tensor.randn(9*27*4).realize()
it = data.cast(dtypes.imagef((9,27,4))).realize()
b1 = it.lazydata.base.realized._buf
del it
it = data.cast(dtypes.imagef((9,27,4))).realize()
assert it.lazydata.base.realized._buf == b1
def test_no_lru_alloc(self):
data = Tensor.randn(9*27*4).realize()
it = data.cast(dtypes.imagef((9,27,4))).realize()
b1 = it.lazydata.base.realized._buf
del it
it = data.cast(dtypes.imagef((10,27,4))).realize()
assert it.lazydata.base.realized._buf != b1
def test_no_lru_alloc_dtype(self):
data = Tensor.randn(9*27*4).realize()
it = data.cast(dtypes.imagef((9,27,4))).realize()
b1 = it.lazydata.base.realized._buf
del it
it = data.cast(dtypes.imageh((9,27,4))).realize()
assert it.lazydata.base.realized._buf != b1
class TestImageIdx(unittest.TestCase):
def test_to_image_idx_real1(self):
gidx0 = Variable('gidx0', 0, 511)
base_idx = (((gidx0*4)%32)*32)+((gidx0//8)%32)
base_valid = gidx0<256
(idx, idy), valid = to_image_idx((4, 64, 4), base_idx, base_valid)
print(idx, idy, idx.min, idx.max, idy.min, idy.max, valid)
assert valid.min == 0
if __name__ == '__main__':
unittest.main()