diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 20 | ||||
| -rw-r--r-- | numpy/core/tests/test_ufunc.py | 19 |
2 files changed, 37 insertions, 2 deletions
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index 6eb0aae55..52299c1be 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -3700,10 +3700,18 @@ PyUFunc_Reduceat(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *ind, /* Copy the first element to start the reduction */ if (otype == NPY_OBJECT) { + /* + * Input (dataptr[0]) and output (dataptr[1]) may + * point to the same memory, e.g. + * np.add.reduceat(a, np.arange(len(a)), out=a). + * In that case need to incref before decref to + * avoid the possibility of the reference count + * being zero temporarily. + */ + Py_XINCREF(*(PyObject **)dataptr_copy[1]); Py_XDECREF(*(PyObject **)dataptr_copy[0]); *(PyObject **)dataptr_copy[0] = *(PyObject **)dataptr_copy[1]; - Py_XINCREF(*(PyObject **)dataptr_copy[0]); } else { memcpy(dataptr_copy[0], dataptr_copy[1], itemsize); @@ -3758,10 +3766,18 @@ PyUFunc_Reduceat(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *ind, /* Copy the first element to start the reduction */ if (otype == NPY_OBJECT) { + /* + * Input (dataptr[0]) and output (dataptr[1]) may point + * to the same memory, e.g. + * np.add.reduceat(a, np.arange(len(a)), out=a). + * In that case need to incref before decref to avoid + * the possibility of the reference count being zero + * temporarily. + */ + Py_XINCREF(*(PyObject **)dataptr_copy[1]); Py_XDECREF(*(PyObject **)dataptr_copy[0]); *(PyObject **)dataptr_copy[0] = *(PyObject **)dataptr_copy[1]; - Py_XINCREF(*(PyObject **)dataptr_copy[0]); } else { memcpy(dataptr_copy[0], dataptr_copy[1], itemsize); diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index ab8cecff0..62fe3c04b 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -665,6 +665,25 @@ class TestUfunc(TestCase): np.add.accumulate(arr, out=arr, axis=-1) assert_array_equal(arr[0, :], np.array([[2]*i for i in [1, 3, 6, 10]])) + def test_object_array_reduceat_inplace(self): + # Checks that in-place reduceats work, see also gh-7465 + arr = np.empty(4, dtype=object) + arr[:] = [[1] for i in range(4)] + out = np.empty(4, dtype=object) + out[:] = [[1] for i in range(4)] + np.add.reduceat(arr, np.arange(4), out=arr) + np.add.reduceat(arr, np.arange(4), out=arr) + assert_array_equal(arr, out) + + # And the same if the axis argument is used + arr = np.ones((2, 4), dtype=object) + arr[0, :] = [[2] for i in range(4)] + out = np.ones((2, 4), dtype=object) + out[0, :] = [[2] for i in range(4)] + np.add.reduceat(arr, np.arange(4), out=arr, axis=-1) + np.add.reduceat(arr, np.arange(4), out=arr, axis=-1) + assert_array_equal(arr, out) + def test_object_scalar_multiply(self): # Tickets #2469 and #4482 arr = np.matrix([1, 2], dtype=object) |