diff options
| -rw-r--r-- | numpy/core/src/multiarray/nditer_api.c | 35 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/nditer_constr.c | 3 | ||||
| -rw-r--r-- | numpy/core/tests/test_nditer.py | 19 |
3 files changed, 37 insertions, 20 deletions
diff --git a/numpy/core/src/multiarray/nditer_api.c b/numpy/core/src/multiarray/nditer_api.c index 81209651b..a1ca5bff5 100644 --- a/numpy/core/src/multiarray/nditer_api.c +++ b/numpy/core/src/multiarray/nditer_api.c @@ -1760,6 +1760,9 @@ npyiter_allocate_buffers(NpyIter *iter, char **errmsg) } goto fail; } + if (PyDataType_FLAGCHK(op_dtype[iop], NPY_NEEDS_INIT)) { + memset(buffer, '\0', itemsize*buffersize); + } buffers[iop] = buffer; } } @@ -2529,16 +2532,18 @@ npyiter_copy_to_buffers(NpyIter *iter, char **prev_dataptrs) skip_transfer = 1; } - /* If the data type requires zero-inititialization */ - if (PyDataType_FLAGCHK(dtypes[iop], NPY_NEEDS_INIT)) { - NPY_IT_DBG_PRINT("Iterator: Buffer requires init, " - "memsetting to 0\n"); - memset(ptrs[iop], 0, dtypes[iop]->elsize*op_transfersize); - /* Can't skip the transfer in this case */ - skip_transfer = 0; - } - - if (!skip_transfer) { + /* + * Copy data to the buffers if necessary. + * + * We always copy if the operand has references. In that case + * a "write" function must be in use that either copies or clears + * the buffer. + * This write from buffer call does not check for skip-transfer + * so we have to assume the buffer is cleared. For dtypes that + * do not have references, we can assume that the write function + * will leave the source (buffer) unmodified. + */ + if (!skip_transfer || PyDataType_REFCHK(dtypes[iop])) { NPY_IT_DBG_PRINT2("Iterator: Copying operand %d to " "buffer (%d items)\n", (int)iop, (int)op_transfersize); @@ -2554,16 +2559,6 @@ npyiter_copy_to_buffers(NpyIter *iter, char **prev_dataptrs) } } } - else if (ptrs[iop] == buffers[iop]) { - /* If the data type requires zero-inititialization */ - if (PyDataType_FLAGCHK(dtypes[iop], NPY_NEEDS_INIT)) { - NPY_IT_DBG_PRINT1("Iterator: Write-only buffer for " - "operand %d requires init, " - "memsetting to 0\n", (int)iop); - memset(ptrs[iop], 0, dtypes[iop]->elsize*transfersize); - } - } - } /* diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c index 982dca849..2197fe798 100644 --- a/numpy/core/src/multiarray/nditer_constr.c +++ b/numpy/core/src/multiarray/nditer_constr.c @@ -594,6 +594,9 @@ NpyIter_Copy(NpyIter *iter) if (buffers[iop] == NULL) { out_of_memory = 1; } + if (PyDataType_FLAGCHK(dtypes[iop], NPY_NEEDS_INIT)) { + memset(buffers[iop], '\0', itemsize*buffersize); + } } } diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py index ddcc8f283..b44343c57 100644 --- a/numpy/core/tests/test_nditer.py +++ b/numpy/core/tests/test_nditer.py @@ -2920,6 +2920,25 @@ def test_object_iter_cleanup(): assert_raises(TypeError, np.logical_or.reduce, np.array([T(), T()], dtype='O')) +def test_object_iter_cleanup_reduce(): + # Similar as above, but a complex reduction case that was previously + # missed (see gh-18810). + # The following array is special in that it cannot be flattened: + arr = np.array([[None, 1], [-1, -1], [None, 2], [-1, -1]])[::2] + with pytest.raises(TypeError): + np.sum(arr) + +@pytest.mark.parametrize("arr", [ + np.ones((8000, 4, 2), dtype=object)[:, ::2, :], + np.ones((8000, 4, 2), dtype=object, order="F")[:, ::2, :], + np.ones((8000, 4, 2), dtype=object)[:, ::2, :].copy("F")]) +def test_object_iter_cleanup_large_reduce(arr): + # More complicated calls are possible for large arrays: + out = np.ones(8000, dtype=np.intp) + # force casting with `dtype=object` + res = np.sum(arr, axis=(1, 2), dtype=object, out=out) + assert_array_equal(res, np.full(8000, 4, dtype=object)) + def test_iter_too_large(): # The total size of the iterator must not exceed the maximum intp due # to broadcasting. Dividing by 1024 will keep it small enough to |