Skip to content

Commit fa3a023

Browse files
mattipcharris
mattip
authored and
charris
committed
BUG: fix for modifying the index arg in ufunc_at
1 parent 5c964bb commit fa3a023

File tree

2 files changed

+79
-33
lines changed

2 files changed

+79
-33
lines changed

‎numpy/core/src/umath/loops.c.src

+72-27
Original file line numberDiff line numberDiff line change
@@ -450,14 +450,19 @@ NPY_NO_EXPORT NPY_GCC_OPT_3 int
450450
void *NPY_UNUSED(func))
451451
{
452452
char *ip1 = args[0];
453-
char *indx = args[1];
453+
char *indxp = args[1];
454454
char *value = args[2];
455455
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
456+
npy_intp shape = steps[3];
456457
npy_intp n = dimensions[0];
457458
npy_intp i;
458459
@type@ *indexed;
459-
for(i = 0; i < n; i++, indx += isindex, value += isb) {
460-
indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx);
460+
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
461+
npy_intp indx = *(npy_intp *)indxp;
462+
if (indx < 0) {
463+
indx += shape;
464+
}
465+
indexed = (@type@ *)(ip1 + is1 * indx);
461466
*indexed = *indexed @OP@ *(@type@ *)value;
462467
}
463468
return 0;
@@ -1241,14 +1246,19 @@ NPY_NO_EXPORT int
12411246
void *NPY_UNUSED(func))
12421247
{
12431248
char *ip1 = args[0];
1244-
char *indx = args[1];
1249+
char *indxp = args[1];
12451250
char *value = args[2];
12461251
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
1252+
npy_intp shape = steps[3];
12471253
npy_intp n = dimensions[0];
12481254
npy_intp i;
12491255
@type@ *indexed;
1250-
for(i = 0; i < n; i++, indx += isindex, value += isb) {
1251-
indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx);
1256+
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
1257+
npy_intp indx = *(npy_intp *)indxp;
1258+
if (indx < 0) {
1259+
indx += shape;
1260+
}
1261+
indexed = (@type@ *)(ip1 + is1 * indx);
12521262
*indexed = npy_floor_divide@c@(*indexed, *(@type@ *)value);
12531263
}
12541264
return 0;
@@ -1395,14 +1405,19 @@ LONGDOUBLE_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
13951405
void *NPY_UNUSED(func))
13961406
{
13971407
char *ip1 = args[0];
1398-
char *indx = args[1];
1408+
char *indxp = args[1];
13991409
char *value = args[2];
14001410
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
1411+
npy_intp shape = steps[3];
14011412
npy_intp n = dimensions[0];
14021413
npy_intp i;
14031414
npy_longdouble *indexed;
1404-
for(i = 0; i < n; i++, indx += isindex, value += isb) {
1405-
indexed = (npy_longdouble *)(ip1 + is1 * *(npy_intp *)indx);
1415+
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
1416+
npy_intp indx = *(npy_intp *)indxp;
1417+
if (indx < 0) {
1418+
indx += shape;
1419+
}
1420+
indexed = (npy_longdouble *)(ip1 + is1 * indx);
14061421
*indexed = *indexed @OP@ *(npy_longdouble *)value;
14071422
}
14081423
return 0;
@@ -1520,14 +1535,19 @@ HALF_@kind@_indexed(void *NPY_UNUSED(context),
15201535
void *NPY_UNUSED(func))
15211536
{
15221537
char *ip1 = args[0];
1523-
char *indx = args[1];
1538+
char *indxp = args[1];
15241539
char *value = args[2];
15251540
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
1541+
npy_intp shape = steps[3];
15261542
npy_intp n = dimensions[0];
15271543
npy_intp i;
15281544
npy_half *indexed;
1529-
for(i = 0; i < n; i++, indx += isindex, value += isb) {
1530-
indexed = (npy_half *)(ip1 + is1 * *(npy_intp *)indx);
1545+
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
1546+
npy_intp indx = *(npy_intp *)indxp;
1547+
if (indx < 0) {
1548+
indx += shape;
1549+
}
1550+
indexed = (npy_half *)(ip1 + is1 * indx);
15311551
const float v = npy_half_to_float(*(npy_half *)value);
15321552
*indexed = npy_float_to_half(npy_half_to_float(*indexed) @OP@ v);
15331553
}
@@ -1641,14 +1661,19 @@ HALF_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
16411661
void *NPY_UNUSED(func))
16421662
{
16431663
char *ip1 = args[0];
1644-
char *indx = args[1];
1664+
char *indxp = args[1];
16451665
char *value = args[2];
16461666
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
1667+
npy_intp shape = steps[3];
16471668
npy_intp n = dimensions[0];
16481669
npy_intp i;
16491670
npy_half *indexed;
1650-
for(i = 0; i < n; i++, indx += isindex, value += isb) {
1651-
indexed = (npy_half *)(ip1 + is1 * *(npy_intp *)indx);
1671+
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
1672+
npy_intp indx = *(npy_intp *)indxp;
1673+
if (indx < 0) {
1674+
indx += shape;
1675+
}
1676+
indexed = (npy_half *)(ip1 + is1 * indx);
16521677
npy_half v = *(npy_half *)value;
16531678
*indexed = (@OP@(*indexed, v) || npy_half_isnan(*indexed)) ? *indexed : v;
16541679
}
@@ -1679,14 +1704,19 @@ HALF_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
16791704
void *NPY_UNUSED(func))
16801705
{
16811706
char *ip1 = args[0];
1682-
char *indx = args[1];
1707+
char *indxp = args[1];
16831708
char *value = args[2];
16841709
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
1710+
npy_intp shape = steps[3];
16851711
npy_intp n = dimensions[0];
16861712
npy_intp i;
16871713
npy_half *indexed;
1688-
for (i = 0; i < n; i++, indx += isindex, value += isb) {
1689-
indexed = (npy_half *)(ip1 + is1 * *(npy_intp *)indx);
1714+
for (i = 0; i < n; i++, indxp += isindex, value += isb) {
1715+
npy_intp indx = *(npy_intp *)indxp;
1716+
if (indx < 0) {
1717+
indx += shape;
1718+
}
1719+
indexed = (npy_half *)(ip1 + is1 * indx);
16901720
npy_half v = *(npy_half *)value;
16911721
*indexed = (@OP@(*indexed, v) || npy_half_isnan(v)) ? *indexed: v;
16921722
}
@@ -1717,14 +1747,19 @@ HALF_floor_divide_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
17171747
void *NPY_UNUSED(func))
17181748
{
17191749
char *ip1 = args[0];
1720-
char *indx = args[1];
1750+
char *indxp = args[1];
17211751
char *value = args[2];
17221752
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
1753+
npy_intp shape = steps[3];
17231754
npy_intp n = dimensions[0];
17241755
npy_intp i;
17251756
npy_half *indexed;
1726-
for(i = 0; i < n; i++, indx += isindex, value += isb) {
1727-
indexed = (npy_half *)(ip1 + is1 * *(npy_intp *)indx);
1757+
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
1758+
npy_intp indx = *(npy_intp *)indxp;
1759+
if (indx < 0) {
1760+
indx += shape;
1761+
}
1762+
indexed = (npy_half *)(ip1 + is1 * indx);
17281763
float v = npy_half_to_float(*(npy_half *)value);
17291764
float div = npy_floor_dividef(npy_half_to_float(*indexed), v);
17301765
*indexed = npy_float_to_half(div);
@@ -1947,14 +1982,19 @@ NPY_NO_EXPORT int @TYPE@_@kind@_indexed
19471982
(PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
19481983
{
19491984
char *ip1 = args[0];
1950-
char *indx = args[1];
1985+
char *indxp = args[1];
19511986
char *value = args[2];
19521987
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
1988+
npy_intp shape = steps[3];
19531989
npy_intp n = dimensions[0];
19541990
npy_intp i;
19551991
@ftype@ *indexed;
1956-
for(i = 0; i < n; i++, indx += isindex, value += isb) {
1957-
indexed = (@ftype@ *)(ip1 + is1 * *(npy_intp *)indx);
1992+
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
1993+
npy_intp indx = *(npy_intp *)indxp;
1994+
if (indx < 0) {
1995+
indx += shape;
1996+
}
1997+
indexed = (@ftype@ *)(ip1 + is1 * indx);
19581998
const @ftype@ b_r = ((@ftype@ *)value)[0];
19591999
const @ftype@ b_i = ((@ftype@ *)value)[1];
19602000
indexed[0] @OP@= b_r;
@@ -1981,14 +2021,19 @@ NPY_NO_EXPORT int @TYPE@_multiply_indexed
19812021
(PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
19822022
{
19832023
char *ip1 = args[0];
1984-
char *indx = args[1];
2024+
char *indxp = args[1];
19852025
char *value = args[2];
19862026
npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
2027+
npy_intp shape = steps[3];
19872028
npy_intp n = dimensions[0];
19882029
npy_intp i;
19892030
@ftype@ *indexed;
1990-
for(i = 0; i < n; i++, indx += isindex, value += isb) {
1991-
indexed = (@ftype@ *)(ip1 + is1 * *(npy_intp *)indx);
2031+
for(i = 0; i < n; i++, indxp += isindex, value += isb) {
2032+
npy_intp indx = *(npy_intp *)indxp;
2033+
if (indx < 0) {
2034+
indx += shape;
2035+
}
2036+
indexed = (@ftype@ *)(ip1 + is1 * indx);
19922037
const @ftype@ a_r = indexed[0];
19932038
const @ftype@ a_i = indexed[1];
19942039
const @ftype@ b_r = ((@ftype@ *)value)[0];

‎numpy/core/src/umath/ufunc_object.c

+7-6
Original file line numberDiff line numberDiff line change
@@ -5939,7 +5939,7 @@ trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
59395939
int buffersize=0, errormask = 0;
59405940
int res;
59415941
char *args[3];
5942-
npy_intp steps[3];
5942+
npy_intp steps[4];
59435943
args[0] = (char *) iter->baseoffset;
59445944
steps[0] = iter->fancy_strides[0];
59455945
if (ufuncimpl->nin == 1) {
@@ -5962,16 +5962,17 @@ trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
59625962
do {
59635963
npy_intp *inner_size = NpyIter_GetInnerLoopSizePtr(iter->outer);
59645964
npy_intp * indxP = (npy_intp *)iter->outer_ptrs[0];
5965-
for (npy_intp i=0; i < *inner_size; i++) {
5966-
if (indxP[i] < 0) {
5967-
indxP[i] += iter->fancy_dims[0];
5968-
}
5969-
}
59705965
args[1] = (char *)indxP;
59715966
steps[1] = iter->outer_strides[0];
5967+
/*
5968+
* The value of iter->fancy_dims[0] is added to negative indexes
5969+
* inside the inner loop
5970+
*/
5971+
steps[3] = iter->fancy_dims[0];
59725972

59735973
res = ufuncimpl->contiguous_indexed_loop(
59745974
context, args, inner_size, steps, NULL);
5975+
59755976
if (args[2] != NULL) {
59765977
args[2] += (*inner_size) * steps[2];
59775978
}

0 commit comments

Comments
 (0)