Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_npz.py 2.41 KB
Newer Older
root's avatar
clone code from github  
root committed
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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100 
import io
import pickle
import unittest

import cupy
from cupy import testing


class TestNpz(unittest.TestCase):

    @testing.for_all_dtypes()
    def test_save_load(self, dtype):
        a = testing.shaped_arange((2, 3, 4), dtype=dtype)
        sio = io.BytesIO()
        cupy.save(sio, a)
        s = sio.getvalue()
        sio.close()

        sio = io.BytesIO(s)
        b = cupy.load(sio)
        sio.close()

        testing.assert_array_equal(a, b)

    def test_save_pickle(self):
        data = object()

        sio = io.BytesIO()
        with self.assertRaises(ValueError):
            cupy.save(sio, data, allow_pickle=False)
        sio.close()

        sio = io.BytesIO()
        cupy.save(sio, data, allow_pickle=True)
        sio.close()

    def test_load_pickle(self):
        a = testing.shaped_arange((2, 3, 4), dtype=cupy.float32)

        sio = io.BytesIO()
        a.dump(sio)
        s = sio.getvalue()
        sio.close()

        sio = io.BytesIO(s)
        b = cupy.load(sio, allow_pickle=True)
        testing.assert_array_equal(a, b)
        sio.close()

        sio = io.BytesIO(s)
        with self.assertRaises(ValueError):
            cupy.load(sio, allow_pickle=False)
        sio.close()

    @testing.for_all_dtypes()
    def check_savez(self, savez, dtype):
        a1 = testing.shaped_arange((2, 3, 4), dtype=dtype)
        a2 = testing.shaped_arange((3, 4, 5), dtype=dtype)

        sio = io.BytesIO()
        savez(sio, a1, a2)
        s = sio.getvalue()
        sio.close()

        sio = io.BytesIO(s)
        with cupy.load(sio) as d:
            b1 = d['arr_0']
            b2 = d['arr_1']
        sio.close()

        testing.assert_array_equal(a1, b1)
        testing.assert_array_equal(a2, b2)

    def test_savez(self):
        self.check_savez(cupy.savez)

    def test_savez_compressed(self):
        self.check_savez(cupy.savez_compressed)

    @testing.for_all_dtypes()
    def test_pickle(self, dtype):
        a = testing.shaped_arange((2, 3, 4), dtype=dtype)
        s = pickle.dumps(a)
        b = pickle.loads(s)
        testing.assert_array_equal(a, b)

    @testing.for_all_dtypes()
    def test_dump(self, dtype):
        a = testing.shaped_arange((2, 3, 4), dtype=dtype)

        sio = io.BytesIO()
        a.dump(sio)
        s = sio.getvalue()
        sio.close()

        sio = io.BytesIO(s)
        b = cupy.load(sio, allow_pickle=True)
        sio.close()

        testing.assert_array_equal(a, b)