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
|
from pypy.module.micronumpy import support
from pypy.module.micronumpy.iterators import ArrayIter
from pypy.module.micronumpy.strides import is_c_contiguous, is_f_contiguous
from pypy.module.micronumpy import constants as NPY
class MockArray(object):
flags = 0
class dtype:
elsize = 1
def __init__(self, shape, strides, start=0):
self.shape = shape
self.strides = strides
self.start = start
if is_c_contiguous(self):
self.flags |= NPY.ARRAY_C_CONTIGUOUS
if is_f_contiguous(self):
self.flags |= NPY.ARRAY_F_CONTIGUOUS
def get_shape(self):
return self.shape
def get_strides(self):
return self.strides
class TestIterDirect(object):
def test_iterator_basic(self):
#Let's get started, simple iteration in C order with
#contiguous layout => strides[-1] is 1
shape = [3, 5]
strides = [5, 1]
backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
assert backstrides == [10, 4]
i = ArrayIter(MockArray(shape, strides), support.product(shape), shape,
strides, backstrides)
assert i.contiguous
s = i.reset()
s = i.next(s)
s = i.next(s)
s = i.next(s)
assert s.offset == 3
assert not i.done(s)
assert s._indices == [0,0]
assert i.indices(s) == [0,3]
#cause a dimension overflow
s = i.next(s)
s = i.next(s)
assert s.offset == 5
assert s._indices == [0,3]
assert i.indices(s) == [1,0]
#Now what happens if the array is transposed? strides[-1] != 1
# therefore layout is non-contiguous
strides = [1, 3]
backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
assert backstrides == [2, 12]
i = ArrayIter(MockArray(shape, strides), support.product(shape), shape,
strides, backstrides)
assert not i.contiguous
s = i.reset()
s = i.next(s)
s = i.next(s)
s = i.next(s)
assert s.offset == 9
assert not i.done(s)
assert s._indices == [0,3]
#cause a dimension overflow
s = i.next(s)
s = i.next(s)
assert s.offset == 1
assert s._indices == [1,0]
def test_one_in_shape(self):
strides = [16, 4, 8]
shape = [3, 4, 1]
backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
assert backstrides == [32, 12, 0]
i = ArrayIter(MockArray(shape, strides), support.product(shape), shape,
strides, backstrides)
assert not i.contiguous
s = i.reset()
for j in range(3):
s = i.next(s)
assert s.offset == 12
assert not i.done(s)
assert s._indices == [0, 3, 0]
while not i.done(s):
old_indices = s._indices[:]
old_offset = s.offset
s = i.next(s)
assert s.offset == 0
assert s._indices == [0, 0, 0]
assert old_indices == [2, 3, 0]
assert old_offset == 44
def test_iterator_goto(self):
shape = [3, 5]
strides = [1, 3]
backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
assert backstrides == [2, 12]
a = MockArray(shape, strides, 42)
i = ArrayIter(a, support.product(shape), shape,
strides, backstrides)
assert not i.contiguous
s = i.reset()
assert s.index == 0
assert s._indices == [0, 0]
assert s.offset == a.start
s = i.goto(11)
assert s.index == 11
assert s._indices is None
assert s.offset == a.start + 5
|
