diff options
| -rw-r--r-- | numpy/core/src/multiarray/nditer_api.c | 3 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/nditer_constr.c | 3 | ||||
| -rw-r--r-- | numpy/core/tests/test_nditer.py | 19 |
3 files changed, 25 insertions, 0 deletions
diff --git a/numpy/core/src/multiarray/nditer_api.c b/numpy/core/src/multiarray/nditer_api.c index 81209651b..063e30919 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; } } 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..c32822944 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 cananot 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 |