Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_iterate.py 4.32 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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138 
import unittest
import warnings

import numpy
import pytest

import cupy
from cupy import testing


class TestFlatiter(unittest.TestCase):

    def test_base(self):
        for xp in (numpy, cupy):
            a = xp.zeros((2, 3, 4))
            assert a.flat.base is a

    def test_iter(self):
        for xp in (numpy, cupy):
            it = xp.zeros((2, 3, 4)).flat
            assert iter(it) is it

    def test_next(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((2, 3, 4), xp)
            e = a.flatten()
            for ai, ei in zip(a.flat, e):
                assert ai == ei

    def test_len(self):
        for xp in (numpy, cupy):
            a = xp.zeros((2, 3, 4))
            assert len(a.flat) == 24
            assert len(a[::2].flat) == 12

    @testing.numpy_cupy_array_equal()
    def test_copy(self, xp):
        a = testing.shaped_arange((2, 3, 4), xp)
        o = a.flat.copy()
        assert a is not o
        return a.flat.copy()

    @testing.numpy_cupy_array_equal()
    def test_copy_next(self, xp):
        a = testing.shaped_arange((2, 3, 4), xp)
        it = a.flat
        it.__next__()
        return it.copy()  # Returns the flattened copy of whole `a`


@testing.parameterize(
    {'shape': (2, 3, 4), 'index': Ellipsis},
    {'shape': (2, 3, 4), 'index': 0},
    {'shape': (2, 3, 4), 'index': 10},
    {'shape': (2, 3, 4), 'index': slice(None)},
    {'shape': (2, 3, 4), 'index': slice(None, 10)},
    {'shape': (2, 3, 4), 'index': slice(None, None, 2)},
    {'shape': (2, 3, 4), 'index': slice(None, None, -1)},
    {'shape': (2, 3, 4), 'index': slice(10, None, -1)},
    {'shape': (2, 3, 4), 'index': slice(10, None, -2)},
    {'shape': (), 'index': slice(None)},
    {'shape': (10,), 'index': slice(None)},
)
class TestFlatiterSubscript(unittest.TestCase):

    @testing.for_CF_orders()
    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_getitem(self, xp, dtype, order):
        a = testing.shaped_arange(self.shape, xp, dtype, order)
        return a.flat[self.index]

    @testing.for_CF_orders()
    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_setitem_scalar(self, xp, dtype, order):
        a = xp.zeros(self.shape, dtype=dtype, order=order)
        a.flat[self.index] = 1
        return a

    @testing.for_CF_orders()
    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_setitem_ndarray_1d(self, xp, dtype, order):
        if numpy.isscalar(self.index):
            pytest.skip()
        a = xp.zeros(self.shape, dtype=dtype, order=order)
        v = testing.shaped_arange((3,), xp, dtype, order)
        a.flat[self.index] = v
        return a

    @testing.for_CF_orders()
    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_setitem_ndarray_nd(self, xp, dtype, order):
        if numpy.isscalar(self.index):
            pytest.skip()
        a = xp.zeros(self.shape, dtype=dtype, order=order)
        v = testing.shaped_arange((2, 3), xp, dtype, order)
        a.flat[self.index] = v
        return a

    @testing.for_CF_orders()
    @testing.for_all_dtypes_combination(('a_dtype', 'v_dtype'))
    @testing.numpy_cupy_array_equal()
    def test_setitem_ndarray_different_types(
            self, xp, a_dtype, v_dtype, order):
        if numpy.isscalar(self.index):
            pytest.skip()
        a = xp.zeros(self.shape, dtype=a_dtype, order=order)
        v = testing.shaped_arange((3,), xp, v_dtype, order)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', numpy.ComplexWarning)
            a.flat[self.index] = v
        return a


@testing.parameterize(
    {'shape': (2, 3, 4), 'index': None},
    {'shape': (2, 3, 4), 'index': (0,)},
    {'shape': (2, 3, 4), 'index': True},
    {'shape': (2, 3, 4), 'index': cupy.array([0])},
    {'shape': (2, 3, 4), 'index': [0]},
)
class TestFlatiterSubscriptIndexError(unittest.TestCase):

    @testing.for_all_dtypes()
    def test_getitem(self, dtype):
        a = testing.shaped_arange(self.shape, cupy, dtype)
        with pytest.raises(IndexError):
            a.flat[self.index]

    @testing.for_all_dtypes()
    def test_setitem(self, dtype):
        a = testing.shaped_arange(self.shape, cupy, dtype)
        v = testing.shaped_arange((1,), cupy, dtype)
        with pytest.raises(IndexError):
            a.flat[self.index] = v