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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
|
import numpy as np
from numpy.testing import *
from numpy.lib.stride_tricks import broadcast_arrays
def assert_shapes_correct(input_shapes, expected_shape):
""" Broadcast a list of arrays with the given input shapes and check the
common output shape.
"""
inarrays = [np.zeros(s) for s in input_shapes]
outarrays = broadcast_arrays(*inarrays)
outshapes = [a.shape for a in outarrays]
expected = [expected_shape] * len(inarrays)
assert outshapes == expected
def assert_incompatible_shapes_raise(input_shapes):
""" Broadcast a list of arrays with the given (incompatible) input shapes
and check that they raise a ValueError.
"""
inarrays = [np.zeros(s) for s in input_shapes]
assert_raises(ValueError, broadcast_arrays, *inarrays)
def assert_same_as_ufunc(shape0, shape1, transposed=False, flipped=False):
""" Broadcast two shapes against each other and check that the data layout
is the same as if a ufunc did the broadcasting.
"""
x0 = np.zeros(shape0, dtype=int)
# Note that multiply.reduce's identity element is 1.0, so when shape1==(),
# this gives the desired n==1.
n = int(np.multiply.reduce(shape1))
x1 = np.arange(n).reshape(shape1)
if transposed:
x0 = x0.T
x1 = x1.T
if flipped:
x0 = x0[::-1]
x1 = x1[::-1]
# Use the add ufunc to do the broadcasting. Since we're adding 0s to x1, the
# result should be exactly the same as the broadcasted view of x1.
y = x0 + x1
b0, b1 = broadcast_arrays(x0, x1)
assert_array_equal(y, b1)
def test_same():
x = np.arange(10)
y = np.arange(10)
bx, by = broadcast_arrays(x, y)
assert_array_equal(x, bx)
assert_array_equal(y, by)
def test_one_off():
x = np.array([[1,2,3]])
y = np.array([[1],[2],[3]])
bx, by = broadcast_arrays(x, y)
bx0 = np.array([[1,2,3],[1,2,3],[1,2,3]])
by0 = bx0.T
assert_array_equal(bx0, bx)
assert_array_equal(by0, by)
def test_same_input_shapes():
""" Check that the final shape is just the input shape.
"""
data = [
(),
(1,),
(3,),
(0,1),
(0,3),
(1,0),
(3,0),
(1,3),
(3,1),
(3,3),
]
for shape in data:
input_shapes = [shape]
# Single input.
yield assert_shapes_correct, input_shapes, shape
# Double input.
input_shapes2 = [shape, shape]
yield assert_shapes_correct, input_shapes2, shape
# Triple input.
input_shapes3 = [shape, shape, shape]
yield assert_shapes_correct, input_shapes3, shape
def test_two_compatible_by_ones_input_shapes():
""" Check that two different input shapes (of the same length but some have
1s) broadcast to the correct shape.
"""
data = [
[[(1,), (3,)], (3,)],
[[(1,3), (3,3)], (3,3)],
[[(3,1), (3,3)], (3,3)],
[[(1,3), (3,1)], (3,3)],
[[(1,1), (3,3)], (3,3)],
[[(1,1), (1,3)], (1,3)],
[[(1,1), (3,1)], (3,1)],
[[(1,0), (0,0)], (0,0)],
[[(0,1), (0,0)], (0,0)],
[[(1,0), (0,1)], (0,0)],
[[(1,1), (0,0)], (0,0)],
[[(1,1), (1,0)], (1,0)],
[[(1,1), (0,1)], (0,1)],
]
for input_shapes, expected_shape in data:
yield assert_shapes_correct, input_shapes, expected_shape
# Reverse the input shapes since broadcasting should be symmetric.
yield assert_shapes_correct, input_shapes[::-1], expected_shape
def test_two_compatible_by_prepending_ones_input_shapes():
""" Check that two different input shapes (of different lengths) broadcast
to the correct shape.
"""
data = [
[[(), (3,)], (3,)],
[[(3,), (3,3)], (3,3)],
[[(3,), (3,1)], (3,3)],
[[(1,), (3,3)], (3,3)],
[[(), (3,3)], (3,3)],
[[(1,1), (3,)], (1,3)],
[[(1,), (3,1)], (3,1)],
[[(1,), (1,3)], (1,3)],
[[(), (1,3)], (1,3)],
[[(), (3,1)], (3,1)],
[[(), (0,)], (0,)],
[[(0,), (0,0)], (0,0)],
[[(0,), (0,1)], (0,0)],
[[(1,), (0,0)], (0,0)],
[[(), (0,0)], (0,0)],
[[(1,1), (0,)], (1,0)],
[[(1,), (0,1)], (0,1)],
[[(1,), (1,0)], (1,0)],
[[(), (1,0)], (1,0)],
[[(), (0,1)], (0,1)],
]
for input_shapes, expected_shape in data:
yield assert_shapes_correct, input_shapes, expected_shape
# Reverse the input shapes since broadcasting should be symmetric.
yield assert_shapes_correct, input_shapes[::-1], expected_shape
def test_incompatible_shapes_raise_valueerror():
""" Check that a ValueError is raised for incompatible shapes.
"""
data = [
[(3,), (4,)],
[(2,3), (2,)],
[(3,), (3,), (4,)],
[(1,3,4), (2,3,3)],
]
for input_shapes in data:
yield assert_incompatible_shapes_raise, input_shapes
# Reverse the input shapes since broadcasting should be symmetric.
yield assert_incompatible_shapes_raise, input_shapes[::-1]
def test_same_as_ufunc():
""" Check that the data layout is the same as if a ufunc did the operation.
"""
data = [
[[(1,), (3,)], (3,)],
[[(1,3), (3,3)], (3,3)],
[[(3,1), (3,3)], (3,3)],
[[(1,3), (3,1)], (3,3)],
[[(1,1), (3,3)], (3,3)],
[[(1,1), (1,3)], (1,3)],
[[(1,1), (3,1)], (3,1)],
[[(1,0), (0,0)], (0,0)],
[[(0,1), (0,0)], (0,0)],
[[(1,0), (0,1)], (0,0)],
[[(1,1), (0,0)], (0,0)],
[[(1,1), (1,0)], (1,0)],
[[(1,1), (0,1)], (0,1)],
[[(), (3,)], (3,)],
[[(3,), (3,3)], (3,3)],
[[(3,), (3,1)], (3,3)],
[[(1,), (3,3)], (3,3)],
[[(), (3,3)], (3,3)],
[[(1,1), (3,)], (1,3)],
[[(1,), (3,1)], (3,1)],
[[(1,), (1,3)], (1,3)],
[[(), (1,3)], (1,3)],
[[(), (3,1)], (3,1)],
[[(), (0,)], (0,)],
[[(0,), (0,0)], (0,0)],
[[(0,), (0,1)], (0,0)],
[[(1,), (0,0)], (0,0)],
[[(), (0,0)], (0,0)],
[[(1,1), (0,)], (1,0)],
[[(1,), (0,1)], (0,1)],
[[(1,), (1,0)], (1,0)],
[[(), (1,0)], (1,0)],
[[(), (0,1)], (0,1)],
]
for input_shapes, expected_shape in data:
yield assert_same_as_ufunc, input_shapes[0], input_shapes[1]
# Reverse the input shapes since broadcasting should be symmetric.
yield assert_same_as_ufunc, input_shapes[1], input_shapes[0]
# Try them transposed, too.
yield assert_same_as_ufunc, input_shapes[0], input_shapes[1], True
# ... and flipped for non-rank-0 inputs in order to test negative
# strides.
if () not in input_shapes:
yield assert_same_as_ufunc, input_shapes[0], input_shapes[1], False, True
yield assert_same_as_ufunc, input_shapes[0], input_shapes[1], True, True
if __name__ == "__main__":
run_module_suite()
|
