Merge pull request #5152 from tovrstra/collapsingbug

juliantaylor · juliantaylor · commit 5a5dc6e8ebba · 2014-10-05T01:04:58.000+02:00
fixed #5147 collapsingbug in einsum
diff --git a/numpy/core/src/multiarray/einsum.c.src b/numpy/core/src/multiarray/einsum.c.src
@@ -2234,7 +2234,7 @@ get_combined_dims_view(PyArrayObject *op, int iop, char *labels)
         else {
             /* Update the combined axis dimensions and strides */
             i = idim + combineoffset;
-            if (combineoffset < 0 &&
+            if (combineoffset < 0 && new_dims[i] != 0 &&
                         new_dims[i] != PyArray_DIM(op, idim)) {
                 PyErr_Format(PyExc_ValueError,
                         "dimensions in operand %d for collapsing "
diff --git a/numpy/core/tests/test_einsum.py b/numpy/core/tests/test_einsum.py
@@ -568,6 +568,16 @@ def test_einsum_fixedstridebug(self):
         tp = np.tensordot(A,B, axes=(0,0))
         assert_equal(es, tp)
 
+    def test_einsum_fixed_collapsingbug(self):
+        # Issue #5147.
+        # The bug only occured when output argument of einssum was used.
+        x = np.random.normal(0, 1, (5, 5, 5, 5))
+        y1 = np.zeros((5, 5))
+        np.einsum('aabb->ab', x, out=y1)
+        idx = np.arange(5)
+        y2 = x[idx[:, None], idx[:, None], idx, idx]
+        assert_equal(y1, y2)
+
 
 if __name__ == "__main__":
     run_module_suite()
