BUG: enforce 'ret is ouf' for 'ret = einsum(..., out=out)'

mattip · mattip · commit 4c7511394578 · 2018-06-19T18:16:29.000-07:00
diff --git a/numpy/core/src/multiarray/einsum.c.src b/numpy/core/src/multiarray/einsum.c.src
@@ -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);
+    if (out) {
+        ret = out;
+    }
+    Py_INCREF(ret);
 
 
     /***************************/
