| 1281 | assert_(c is c2) |
| 1282 | |
| 1283 | def test_iter_op_axes(): |
| 1284 | # Check that custom axes work |
| 1285 | |
| 1286 | # Reverse the axes |
| 1287 | a = arange(6).reshape(2, 3) |
| 1288 | i = nditer([a, a.T], [], [['readonly']]*2, op_axes=[[0, 1], [1, 0]]) |
| 1289 | assert_(all([x == y for (x, y) in i])) |
| 1290 | a = arange(24).reshape(2, 3, 4) |
| 1291 | i = nditer([a.T, a], [], [['readonly']]*2, op_axes=[[2, 1, 0], None]) |
| 1292 | assert_(all([x == y for (x, y) in i])) |
| 1293 | |
| 1294 | # Broadcast 1D to any dimension |
| 1295 | a = arange(1, 31).reshape(2, 3, 5) |
| 1296 | b = arange(1, 3) |
| 1297 | i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [0, -1, -1]]) |
| 1298 | assert_equal([x*y for (x, y) in i], (a*b.reshape(2, 1, 1)).ravel()) |
| 1299 | b = arange(1, 4) |
| 1300 | i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [-1, 0, -1]]) |
| 1301 | assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 3, 1)).ravel()) |
| 1302 | b = arange(1, 6) |
| 1303 | i = nditer([a, b], [], [['readonly']]*2, |
| 1304 | op_axes=[None, [np.newaxis, np.newaxis, 0]]) |
| 1305 | assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 1, 5)).ravel()) |
| 1306 | |
| 1307 | # Inner product-style broadcasting |
| 1308 | a = arange(24).reshape(2, 3, 4) |
| 1309 | b = arange(40).reshape(5, 2, 4) |
| 1310 | i = nditer([a, b], ['multi_index'], [['readonly']]*2, |
| 1311 | op_axes=[[0, 1, -1, -1], [-1, -1, 0, 1]]) |
| 1312 | assert_equal(i.shape, (2, 3, 5, 2)) |
| 1313 | |
| 1314 | # Matrix product-style broadcasting |
| 1315 | a = arange(12).reshape(3, 4) |
| 1316 | b = arange(20).reshape(4, 5) |
| 1317 | i = nditer([a, b], ['multi_index'], [['readonly']]*2, |
| 1318 | op_axes=[[0, -1], [-1, 1]]) |
| 1319 | assert_equal(i.shape, (3, 5)) |
| 1320 | |
| 1321 | def test_iter_op_axes_errors(): |
| 1322 | # Check that custom axes throws errors for bad inputs |