Skip to content

Commit

Permalink
Merge pull request numpy#24294 from charris/backport-24272
Browse files Browse the repository at this point in the history 
BUG: do not modify the input to ufunc_at
  • Loading branch information
@charris
charris committed Jul 30, 2023
2 parents eb20535 + fa3a023 commit f92167b
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 33 deletions.
99 changes: 72 additions & 27 deletions numpy/core/src/umath/loops.c.src
View file
Original file line number Diff line number Diff line change
Expand Up @@ -450,14 +450,19 @@ NPY_NO_EXPORT NPY_GCC_OPT_3 int
void *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 @@ -1241,14 +1246,19 @@ NPY_NO_EXPORT int
void *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 = npy_floor_divide@c@(*indexed, *(@type@ *)value);
}
return 0;
Expand Down Expand Up @@ -1395,14 +1405,19 @@ LONGDOUBLE_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
void *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;
npy_longdouble *indexed;
for(i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (npy_longdouble *)(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 = (npy_longdouble *)(ip1 + is1 * indx);
*indexed = *indexed @OP@ *(npy_longdouble *)value;
}
return 0;
Expand Down Expand Up @@ -1520,14 +1535,19 @@ HALF_@kind@_indexed(void *NPY_UNUSED(context),
void *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;
npy_half *indexed;
for(i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (npy_half *)(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 = (npy_half *)(ip1 + is1 * indx);
const float v = npy_half_to_float(*(npy_half *)value);
*indexed = npy_float_to_half(npy_half_to_float(*indexed) @OP@ v);
}
Expand Down Expand Up @@ -1641,14 +1661,19 @@ HALF_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
void *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;
npy_half *indexed;
for(i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (npy_half *)(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 = (npy_half *)(ip1 + is1 * indx);
npy_half v = *(npy_half *)value;
*indexed = (@OP@(*indexed, v) || npy_half_isnan(*indexed)) ? *indexed : v;
}
Expand Down Expand Up @@ -1679,14 +1704,19 @@ HALF_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
void *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;
npy_half *indexed;
for (i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (npy_half *)(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 = (npy_half *)(ip1 + is1 * indx);
npy_half v = *(npy_half *)value;
*indexed = (@OP@(*indexed, v) || npy_half_isnan(v)) ? *indexed: v;
}
Expand Down Expand Up @@ -1717,14 +1747,19 @@ HALF_floor_divide_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
void *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;
npy_half *indexed;
for(i = 0; i < n; i++, indx += isindex, value += isb) {
indexed = (npy_half *)(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 = (npy_half *)(ip1 + is1 * indx);
float v = npy_half_to_float(*(npy_half *)value);
float div = npy_floor_dividef(npy_half_to_float(*indexed), v);
*indexed = npy_float_to_half(div);
Expand Down Expand Up @@ -1947,14 +1982,19 @@ NPY_NO_EXPORT int @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];
indexed[0] @OP@= b_r;
Expand All @@ -1981,14 +2021,19 @@ NPY_NO_EXPORT int @TYPE@_multiply_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@ a_r = indexed[0];
const @ftype@ a_i = indexed[1];
const @ftype@ b_r = ((@ftype@ *)value)[0];
Expand Down
13 changes: 7 additions & 6 deletions numpy/core/src/umath/ufunc_object.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5939,7 +5939,7 @@ trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
int buffersize=0, errormask = 0;
int res;
char *args[3];
npy_intp steps[3];
npy_intp steps[4];
args[0] = (char *) iter->baseoffset;
steps[0] = iter->fancy_strides[0];
if (ufuncimpl->nin == 1) {
Expand All @@ -5962,16 +5962,17 @@ trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
do {
npy_intp *inner_size = NpyIter_GetInnerLoopSizePtr(iter->outer);
npy_intp * indxP = (npy_intp *)iter->outer_ptrs[0];
for (npy_intp i=0; i < *inner_size; i++) {
if (indxP[i] < 0) {
indxP[i] += iter->fancy_dims[0];
}
}
args[1] = (char *)indxP;
steps[1] = iter->outer_strides[0];
/*
* The value of iter->fancy_dims[0] is added to negative indexes
* inside the inner loop
*/
steps[3] = iter->fancy_dims[0];

res = ufuncimpl->contiguous_indexed_loop(
context, args, inner_size, steps, NULL);

if (args[2] != NULL) {
args[2] += (*inner_size) * steps[2];
}
Expand Down
1 change: 1 addition & 0 deletions numpy/core/tests/test_ufunc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2222,6 +2222,7 @@ def test_at_negative_indexes(self):
assert a[-1] == 11 # issue 24147
assert a[1] == 2
assert a[2] == 3
assert np.all(indxs == [-1, 1, -1, 2])

def test_at_not_none_signature(self):
# Test ufuncs with non-trivial signature raise a TypeError
Expand Down

0 comments on commit f92167b

Please sign in to comment.