| 2626 | assert_equal(it2.operands[1].sum(), a.size) |
| 2627 | |
| 2628 | def test_iter_buffering_reduction(): |
| 2629 | # Test doing buffered reductions with the iterator |
| 2630 | |
| 2631 | a = np.arange(6) |
| 2632 | b = np.array(0., dtype='f8').byteswap().newbyteorder() |
| 2633 | i = nditer([a, b], ['reduce_ok', 'buffered'], |
| 2634 | [['readonly'], ['readwrite', 'nbo']], |
| 2635 | op_axes=[[0], [-1]]) |
| 2636 | with i: |
| 2637 | assert_equal(i[1].dtype, np.dtype('f8')) |
| 2638 | assert_(i[1].dtype != b.dtype) |
| 2639 | # Do the reduction |
| 2640 | for x, y in i: |
| 2641 | y[...] += x |
| 2642 | # Since no axes were specified, should have allocated a scalar |
| 2643 | assert_equal(b, np.sum(a)) |
| 2644 | |
| 2645 | a = np.arange(6).reshape(2, 3) |
| 2646 | b = np.array([0, 0], dtype='f8').byteswap().newbyteorder() |
| 2647 | i = nditer([a, b], ['reduce_ok', 'external_loop', 'buffered'], |
| 2648 | [['readonly'], ['readwrite', 'nbo']], |
| 2649 | op_axes=[[0, 1], [0, -1]]) |
| 2650 | # Reduction shape/strides for the output |
| 2651 | with i: |
| 2652 | assert_equal(i[1].shape, (3,)) |
| 2653 | assert_equal(i[1].strides, (0,)) |
| 2654 | # Do the reduction |
| 2655 | for x, y in i: |
| 2656 | # Use a for loop instead of ``y[...] += x`` |
| 2657 | # (equivalent to ``y[...] = y[...].copy() + x``), |
| 2658 | # because y has zero strides we use for the reduction |
| 2659 | for j in range(len(y)): |
| 2660 | y[j] += x[j] |
| 2661 | assert_equal(b, np.sum(a, axis=1)) |
| 2662 | |
| 2663 | # Iterator inner double loop was wrong on this one |
| 2664 | p = np.arange(2) + 1 |
| 2665 | it = np.nditer([p, None], |
| 2666 | ['delay_bufalloc', 'reduce_ok', 'buffered', 'external_loop'], |
| 2667 | [['readonly'], ['readwrite', 'allocate']], |
| 2668 | op_axes=[[-1, 0], [-1, -1]], |
| 2669 | itershape=(2, 2)) |
| 2670 | with it: |
| 2671 | it.operands[1].fill(0) |
| 2672 | it.reset() |
| 2673 | assert_equal(it[0], [1, 2, 1, 2]) |
| 2674 | |
| 2675 | # Iterator inner loop should take argument contiguity into account |
| 2676 | x = np.ones((7, 13, 8), np.int8)[4:6,1:11:6,1:5].transpose(1, 2, 0) |
| 2677 | x[...] = np.arange(x.size).reshape(x.shape) |
| 2678 | y_base = np.arange(4*4, dtype=np.int8).reshape(4, 4) |
| 2679 | y_base_copy = y_base.copy() |
| 2680 | y = y_base[::2,:,None] |
| 2681 | |
| 2682 | it = np.nditer([y, x], |
| 2683 | ['buffered', 'external_loop', 'reduce_ok'], |
| 2684 | [['readwrite'], ['readonly']]) |
| 2685 | with it: |