diff options
| -rw-r--r-- | numpy/core/src/multiarray/einsum.c.src | 2 | ||||
| -rw-r--r-- | numpy/core/tests/test_einsum.py | 10 |
2 files changed, 11 insertions, 1 deletions
diff --git a/numpy/core/src/multiarray/einsum.c.src b/numpy/core/src/multiarray/einsum.c.src index b3148f573..3f47ec1f4 100644 --- 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 index 226bde0a3..8ac4f2213 100644 --- a/numpy/core/tests/test_einsum.py +++ b/numpy/core/tests/test_einsum.py @@ -568,6 +568,16 @@ class TestEinSum(TestCase): 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() |