MCPcopy
hub / github.com/numpy/numpy / test_iter_op_axes

Function test_iter_op_axes

numpy/_core/tests/test_nditer.py:1421–1457  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

1419 assert_(c is c2)
1420
1421def test_iter_op_axes():
1422 # Check that custom axes work
1423
1424 # Reverse the axes
1425 a = arange(6).reshape(2, 3)
1426 i = nditer([a, a.T], [], [['readonly']] * 2, op_axes=[[0, 1], [1, 0]])
1427 assert_(all([x == y for (x, y) in i]))
1428 a = arange(24).reshape(2, 3, 4)
1429 i = nditer([a.T, a], [], [['readonly']] * 2, op_axes=[[2, 1, 0], None])
1430 assert_(all([x == y for (x, y) in i]))
1431
1432 # Broadcast 1D to any dimension
1433 a = arange(1, 31).reshape(2, 3, 5)
1434 b = arange(1, 3)
1435 i = nditer([a, b], [], [['readonly']] * 2, op_axes=[None, [0, -1, -1]])
1436 assert_equal([x * y for (x, y) in i], (a * b.reshape(2, 1, 1)).ravel())
1437 b = arange(1, 4)
1438 i = nditer([a, b], [], [['readonly']] * 2, op_axes=[None, [-1, 0, -1]])
1439 assert_equal([x * y for (x, y) in i], (a * b.reshape(1, 3, 1)).ravel())
1440 b = arange(1, 6)
1441 i = nditer([a, b], [], [['readonly']] * 2,
1442 op_axes=[None, [np.newaxis, np.newaxis, 0]])
1443 assert_equal([x * y for (x, y) in i], (a * b.reshape(1, 1, 5)).ravel())
1444
1445 # Inner product-style broadcasting
1446 a = arange(24).reshape(2, 3, 4)
1447 b = arange(40).reshape(5, 2, 4)
1448 i = nditer([a, b], ['multi_index'], [['readonly']] * 2,
1449 op_axes=[[0, 1, -1, -1], [-1, -1, 0, 1]])
1450 assert_equal(i.shape, (2, 3, 5, 2))
1451
1452 # Matrix product-style broadcasting
1453 a = arange(12).reshape(3, 4)
1454 b = arange(20).reshape(4, 5)
1455 i = nditer([a, b], ['multi_index'], [['readonly']] * 2,
1456 op_axes=[[0, -1], [-1, 1]])
1457 assert_equal(i.shape, (3, 5))
1458
1459def test_iter_op_axes_errors():
1460 # Check that custom axes throws errors for bad inputs

Callers

nothing calls this directly

Calls 5

assert_Function · 0.90
allFunction · 0.90
assert_equalFunction · 0.90
reshapeMethod · 0.80
ravelMethod · 0.45

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…