8000 BUG: Ensure out is returned in einsum. by mattip · Pull Request #11406 · numpy/numpy · GitHub
[go: up one dir, main page]
Skip to content

BUG: Ensure out is returned in einsum. #11406

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jun 29, 2018
Merged

BUG: Ensure out is returned in einsum. #11406

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d829dc5
BUG: ensure ret is out in einsum
mattip Jun 21, 2018
4f5ada2
check for unlikely error in Assign_Zero
mattip Jun 24, 2018
ae5000f
MAINT: cleanup ret assignment
mattip Jun 27, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 10 additions & 18 deletions numpy/core/src/multiarray/einsum.c.src
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2767,11 +2767,11 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
goto fail;
}

/* Initialize the output to all zeros and reset the iterator */
/* Initialize the output to all zeros */
ret = NpyIter_GetOperandArray(iter)[nop];
Py_INCREF(ret);
PyArray_AssignZero(ret, NULL);

if (PyArray_AssignZero(ret, NULL) < 0) {
goto fail;
}

/***************************/
/*
Expand All @@ -2785,16 +2785,12 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
case 1:
if (ndim == 2) {
if (unbuffered_loop_nop1_ndim2(iter) < 0) {
Py_DECREF(ret);
ret = NULL;
goto fail;
}
goto finish;
}
else if (ndim == 3) {
if (unbuffered_loop_nop1_ndim3(iter) < 0) {
Py_DECREF(ret);
ret = NULL;
goto fail;
}
goto finish;
Expand All @@ -2803,16 +2799,12 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
case 2:
if (ndim == 2) {
if (unbuffered_loop_nop2_ndim2(iter) < 0) {
Py_DECREF(ret);
ret = NULL;
goto fail;
}
goto finish;
}
else if (ndim == 3) {
if (unbuffered_loop_nop2_ndim3(iter) < 0) {
Py_DECREF(ret);
ret = NULL;
goto fail;
}
goto finish;
Expand All @@ -2823,7 +2815,6 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
/***************************/

if (NpyIter_Reset(iter, NULL) != NPY_SUCCEED) {
Py_DECREF(ret);
goto fail;
}

Expand All @@ -2845,8 +2836,6 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
if (sop == NULL) {
PyErr_SetString(PyExc_TypeError,
"invalid data type for einsum");
Py_DECREF(ret);
ret = NULL;
}
else if (NpyIter_GetIterSize(iter) != 0) {
NpyIter_IterNextFunc *iternext;
Expand All @@ -2858,7 +2847,6 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
iternext = NpyIter_GetIterNext(iter, NULL);
if (iternext == NULL) {
NpyIter_Deallocate(iter);
Py_DECREF(ret);
goto fail;
}
dataptr = NpyIter_GetDataPtrArray(iter);
Expand All @@ -2874,12 +2862,16 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,

/* If the API was needed, it may have thrown an error */
if (NpyIter_IterationNeedsAPI(iter) && PyErr_Occurred()) {
Py_DECREF(ret);
ret = NULL;
goto fail;
}
}

finish:
if (out != NULL) {
ret = out;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Out of curiousity, a couple of questions.

  • The iterator iterates over out if present?
  • The writeback to out or the new output array happens in iter deallocation?
  • I assume deallocation also decrefs the out or new array?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I need to incref the output before NpyIterDeallocate so it stays alive through that call.

  • The assignment op[nop] = out near the beginning ensures out, if present is used in iteration, possibly with writeback semantics. Note that nop is the number of inputs, but actually op uses nop +1 entries, one additional for the output.
  • The writeback happens in NpyIterDeallocate, which also decrefs all operands including the output (which were increfed in NpyIter_AdvancedNew.
  • In addition, einsum itself increfed all the inputs while deciding whether to create views or not, around get_combined_dims_view. So whether fail or finish, at the end the function decrefs the inputs, (not the output).

}
Py_INCREF(ret);

NpyIter_Deallocate(iter);
for (iop = 0; iop < nop; ++iop) {
Py_DECREF(op[iop]);
Expand Down
5 changes: 5 additions & 0 deletions numpy/core/tests/test_einsum.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -730,6 +730,11 @@ def test_small_boolean_arrays(self):
res = np.einsum('...ij,...jk->...ik', a, a, out=out)
assert_equal(res, tgt)

def test_out_is_res(self):
a = np.arange(9).reshape(3, 3)
res = np.einsum('...ij,...jk->...ik', a, a, out=a)
assert res is a

def optimize_compare(self, string):
# Tests all paths of the optimization function against
# conventional einsum
Expand Down
0