test_nditer.py 文件源码

def test_iter_nested_iters_broadcast():
    # Test nested iteration with broadcasting
    a = arange(2).reshape(2, 1)
    b = arange(3).reshape(1, 3)

    i, j = np.nested_iters([a, b], [[0], [1]])
    vals = []
    for x in i:
        vals.append([y for y in j])
    assert_equal(vals, [[[0, 0], [0, 1], [0, 2]], [[1, 0], [1, 1], [1, 2]]])

    i, j = np.nested_iters([a, b], [[1], [0]])
    vals = []
    for x in i:
        vals.append([y for y in j])
    assert_equal(vals, [[[0, 0], [1, 0]], [[0, 1], [1, 1]], [[0, 2], [1, 2]]])