8000 refactor ufunc_at to first get the inner loop, then create the iterators · mattip/numpy@3ba259d · GitHub
[go: up one dir, main page]

Skip to content

Commit 3ba259d

Browse files
committed
refactor ufunc_at to first get the inner loop, then create the iterators
1 parent 90243f5 commit 3ba259d

File tree

1 file changed

+40
-41
lines changed

1 file changed

+40
-41
lines changed

numpy/core/src/umath/ufunc_object.c

Lines changed: 40 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -6238,49 +6238,11 @@ ufunc_at(PyUFuncObject *ufunc, PyObject *args)
62386238
}
62396239
}
62406240

6241-
/*
6242-
* Create a map iterator (there is no multi-mapiter, so we need at least 2
6243-
* iterators: one for the mapping, and another for the second operand)
6244-
*/
6245-
iter = (PyArrayMapIterObject *)PyArray_MapIterArrayCopyIfOverlap(
6246-
op1_array, idx, 1, op2_array);
6247-
if (iter == NULL) {
6248-
goto fail;
6249-
}
6250-
op1_array = iter->array; /* May be updateifcopied on overlap */
6251-
6252-
if (op2_array != NULL) {
6253-
/*
6254-
* May need to swap axes so that second operand is
6255-
* iterated over correctly
6256-
*/
6257-
if ((iter->subspace != NULL) && (iter->consec)) {
6258-
PyArray_MapIterSwapAxes(iter, &op2_array, 0);
6259-
if (op2_array == NULL) {
6260-
/* only on memory allocation failure */
6261-
goto fail;
6262-
}
6263-
}
6264-
6265-
/*
6266-
* Create array iter object for second operand that
6267-
* "matches" the map iter object for the first operand.
6268-
* Then we can just iterate over the first and second
6269-
* operands at the same time and not have to worry about
6270-
* picking the correct elements from each operand to apply
6271-
* the ufunc to.
6272-
*/
6273-
if ((iter2 = (PyArrayIterObject *)\
6274-
PyArray_BroadcastToShape((PyObject *)op2_array,
6275-
iter->dimensions, iter->nd))==NULL) {
6276-
goto fail;
6277-
}
6278-
}
6279-
62806241
PyArrayMethodObject *ufuncimpl = NULL;
6281-
62826242
{
62836243
/* Do all the dtype handling and find the correct ufuncimpl */
6244+
Py_INCREF(PyArray_DESCR(op1_array));
6245+
62846246

62856247
PyArrayObject *tmp_operands[3] = {NULL, NULL, NULL};
62866248
PyArray_DTypeMeta *signature[3] = {NULL, NULL, NULL};
@@ -6343,7 +6305,44 @@ ufunc_at(PyUFuncObject *ufunc, PyObject *args)
63436305
}
63446306
}
63456307

6346-
Py_INCREF(PyArray_DESCR(op1_array));
6308+
/*
6309+
* Create a map iterator (there is no multi-mapiter, so we need at least 2
6310+
* iterators: one for the mapping, and another for the second operand)
6311+
*/
6312+
iter = (PyArrayMapIterObject *)PyArray_MapIterArrayCopyIfOverlap(
6313+
op1_array, idx, 1, op2_array);
6314+
if (iter == NULL) {
6315+
goto fail;
6316+
}
6317+
op1_array = iter->array; /* May be updateifcopied on overlap */
6318+
6319+
if (op2_array != NULL) {
6320+
/*
6321+
* May need to swap axes so that second operand is
6322+
* iterated over correctly
6323+
*/
6324+
if ((iter->subspace != NULL) && (iter->consec)) {
6325+
PyArray_MapIterSwapAxes(iter, &op2_array, 0);
6326+
if (op2_array == NULL) {
6327+
/* only on memory allocation failure */
6328+
goto fail;
6329+
}
6330+
}
6331+
6332+
/*
6333+
* Create array iter object for second operand that
6334+
* "matches" the map iter object for the first operand.
6335+
* Then we can just iterate over the first and second
6336+
* operands at the same time and not have to worry about
6337+
* picking the correct elements from each operand to apply
6338+
* the ufunc to.
6339+
*/
6340+
if ((iter2 = (PyArrayIterObject *)\
6341+
PyArray_BroadcastToShape((PyObject *)op2_array,
6342+
iter->dimensions, iter->nd))==NULL) {
6343+
goto fail;
6344+
}
6345+
}
63476346

63486347
PyArrayMethod_Context context = {
63496348
.caller = (PyObject *)ufunc,

0 commit comments

Comments
 (0)
0