Skip to content

Commit

Permalink
Merge pull request numpy#24295 from charris/backport-24281
Browse files Browse the repository at this point in the history
BUG: Further fixes to indexing loop and added tests
  • Loading branch information
charris committed Jul 30, 2023
2 parents f92167b + 7b5e51a commit 9941dfe
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 25 deletions.
11 changes: 8 additions & 3 deletions numpy/core/src/umath/_umath_tests.c.src
Original file line number Diff line number Diff line change
Expand Up @@ -375,13 +375,18 @@ INT32_negative_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
{
char *ip1 = args[0];
char *indx = args[1];
char *indxp = args[1];
npy_intp is1 = steps[0], isindex = steps[1];
npy_intp n = dimensions[0];
npy_intp shape = steps[3];
npy_intp i;
int32_t *indexed;
for(i = 0; i < n; i++, indx += isindex) {
indexed = (int32_t *)(ip1 + is1 * *(npy_intp *)indx);
for(i = 0; i < n; i++, indxp += isindex) {
npy_intp indx = *(npy_intp *)indxp;
if (indx < 0) {
indx += shape;
}
indexed = (int32_t *)(ip1 + is1 * indx);
if (i == 3) {
*indexed = -200;
} else {
Expand Down
22 changes: 16 additions & 6 deletions numpy/core/src/umath/loops_arithm_fp.dispatch.c.src
Original file line number Diff line number Diff line change
Expand Up @@ -258,14 +258,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@_indexed)
(PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
{
char *ip1 = args[0];
char *indx = args[1];
char *indxp = args[1];
char *value = args[2];
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
npy_intp shape = steps[3];
npy_intp n = dimensions[0];
npy_intp i;
@type@ *indexed;
for(i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx);
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
npy_intp indx = *(npy_intp *)indxp;
if (indx < 0) {
indx += shape;
}
indexed = (@type@ *)(ip1 + is1 * indx);
*indexed = *indexed @OP@ *(@type@ *)value;
}
return 0;
Expand Down Expand Up @@ -650,14 +655,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@_indexed)
(PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
{
char *ip1 = args[0];
char *indx = args[1];
char *indxp = args[1];
char *value = args[2];
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
npy_intp shape = steps[3];
npy_intp n = dimensions[0];
npy_intp i;
@ftype@ *indexed;
for(i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (@ftype@ *)(ip1 + is1 * *(npy_intp *)indx);
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
npy_intp indx = *(npy_intp *)indxp;
if (indx < 0) {
indx += shape;
}
indexed = (@ftype@ *)(ip1 + is1 * indx);
const @ftype@ b_r = ((@ftype@ *)value)[0];
const @ftype@ b_i = ((@ftype@ *)value)[1];
#if @is_mul@
Expand Down
22 changes: 16 additions & 6 deletions numpy/core/src/umath/loops_arithmetic.dispatch.c.src
Original file line number Diff line number Diff line change
Expand Up @@ -400,14 +400,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_divide_indexed)
(PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
{
char *ip1 = args[0];
char *indx = args[1];
char *indxp = args[1];
char *value = args[2];
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
npy_intp shape = steps[3];
npy_intp n = dimensions[0];
npy_intp i;
@type@ *indexed;
for(i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx);
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
npy_intp indx = *(npy_intp *)indxp;
if (indx < 0) {
indx += shape;
}
indexed = (@type@ *)(ip1 + is1 * indx);
*indexed = floor_div_@TYPE@(*indexed, *(@type@ *)value);
}
return 0;
Expand Down Expand Up @@ -486,14 +491,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_divide_indexed)
(PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
{
char *ip1 = args[0];
char *indx = args[1];
char *indxp = args[1];
char *value = args[2];
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
npy_intp shape = steps[3];
npy_intp n = dimensions[0];
npy_intp i;
@type@ *indexed;
for(i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx);
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
npy_intp indx = *(npy_intp *)indxp;
if (indx < 0) {
indx += shape;
}
indexed = (@type@ *)(ip1 + is1 * indx);
@type@ in2 = *(@type@ *)value;
if (NPY_UNLIKELY(in2 == 0)) {
npy_set_floatstatus_divbyzero();
Expand Down
11 changes: 8 additions & 3 deletions numpy/core/src/umath/loops_minmax.dispatch.c.src
Original file line number Diff line number Diff line change
Expand Up @@ -456,14 +456,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@_indexed)
(PyArrayMethod_Context *NPY_UNUSED(context), char *const *args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
{
char *ip1 = args[0];
char *indx = args[1];
char *indxp = args[1];
char *value = args[2];
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
npy_intp n = dimensions[0];
npy_intp shape = steps[3];
npy_intp i;
@type@ *indexed;
for(i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx);
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
npy_intp indx = *(npy_intp *)indxp;
if (indx < 0) {
indx += shape;
}
indexed = (@type@ *)(ip1 + is1 * indx);
*indexed = SCALAR_OP(*indexed, *(@type@ *)value);
}
return 0;
Expand Down
22 changes: 15 additions & 7 deletions numpy/core/tests/test_ufunc.py
Original file line number Diff line number Diff line change
Expand Up @@ -2215,13 +2215,21 @@ def test_ufunc_at_advanced(self):
np.maximum.at(a, [0], 0)
assert_equal(a, np.array([1, 2, 3]))

def test_at_negative_indexes(self):
a = np.arange(10)
indxs = np.array([-1, 1, -1, 2])
np.add.at(a, indxs, 1)
assert a[-1] == 11 # issue 24147
assert a[1] == 2
assert a[2] == 3
@pytest.mark.parametrize("dtype",
np.typecodes['AllInteger'] + np.typecodes['Float'])
@pytest.mark.parametrize("ufunc",
[np.add, np.subtract, np.divide, np.minimum, np.maximum])
def test_at_negative_indexes(self, dtype, ufunc):
a = np.arange(0, 10).astype(dtype)
indxs = np.array([-1, 1, -1, 2]).astype(np.intp)
vals = np.array([1, 5, 2, 10], dtype=a.dtype)

expected = a.copy()
for i, v in zip(indxs, vals):
expected[i] = ufunc(expected[i], v)

ufunc.at(a, indxs, vals)
assert_array_equal(a, expected)
assert np.all(indxs == [-1, 1, -1, 2])

def test_at_not_none_signature(self):
Expand Down

0 comments on commit 9941dfe

Please sign in to comment.