BUG: Ensure out is returned in einsum. #11406
Conversation
|
Is this in #11345? |
|
IMO it is orthogonal to #11345 |
| @@ -2769,8 +2769,11 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop, | |||
|
|
|||
| /* Initialize the output to all zeros and reset the iterator */ | |||
| ret = NpyIter_GetOperandArray(iter)[nop]; | |||
| Py_INCREF(ret); | |||
| PyArray_AssignZero(ret, NULL); | |||
There was a problem hiding this comment.
Can you convince me why this line should before the if and not after it?
There was a problem hiding this comment.
Doing PyArray_AssignZero on out will fail since out, with activated writeback semantics, will have its READONLY flag set. We must only write to the operand ndarray, the one whose base is out and has extra memory allocated for the writeback buffer. ret is subsequently reassigned to out since that is what must be returned after wrtieback resolution which resets the READONLY flag to its previous value.
96cd478 to
4f5ada2
Compare
|
CircleCI failed and I cannot restart the job, closing and reopening to trip CI |
| if (PyArray_AssignZero(ret, NULL) < 0) { | ||
| goto fail; | ||
| } | ||
| if (out != NULL) { |
There was a problem hiding this comment.
This looks fishy, why do anything to ret if it is simply set to out?
There was a problem hiding this comment.
If the output operand has writeback semantics, NpyIter_GetOperandArray(iter)[nop] is the temporary ndarray created with out as it's base ndarray. Then we zero the temporary data rather than changing out itself. See issue #10956 for the python-level equivalent discussion about what is meant by iter.operand[i] - the temporary ndarray with writeback semantics or the underlying user-created input-ouput ndarray.
Lines 2775-2777 ensure that ret and out point to the same ndarray, as required when using an out kwarg.
|
Can you move the |
|
Moving |
| } | ||
| } | ||
|
|
||
| finish: | ||
| if (out != NULL) { | ||
| ret = out; |
There was a problem hiding this comment.
Out of curiousity, a couple of questions.
- The iterator iterates over
outif present? - The writeback to
outor the new output array happens in iter deallocation? - I assume deallocation also decrefs the
outor new array?
There was a problem hiding this comment.
I need to incref the output before NpyIterDeallocate so it stays alive through that call.
- The assignment
op[nop] = outnear the beginning ensuresout, if present is used in iteration, possibly with writeback semantics. Note thatnopis the number of inputs, but actuallyopusesnop +1entries, one additional for the output. - The writeback happens in
NpyIterDeallocate, which also decrefs all operands including the output (which were increfed inNpyIter_AdvancedNew. - In addition,
einsumitself increfed all the inputs while deciding whether to create views or not, aroundget_combined_dims_view. So whether fail or finish, at the end the function decrefs the inputs, (not the output).
|
Thanks Matti. |
In einsum, if out is given, it should be assigned to ret.