8000 Merge pull request #5152 from tovrstra/collapsingbug · numpy/numpy@5a5dc6e · GitHub
[go: up one dir, main page]

Skip to content

Commit 5a5dc6e

Browse files
committed
Merge pull request #5152 from tovrstra/collapsingbug
fixed #5147 collapsingbug in einsum
2 parents ed228e6 + 8615d5e commit 5a5dc6e

File tree

2 files changed

+11
-1
lines changed

2 files changed

+11
-1
lines changed

numpy/core/src/multiarray/einsum.c.src

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2234,7 +2234,7 @@ get_combined_dims_view(PyArrayObject *op, int iop, char *labels)
22342234
else {
22352235
/* Update the combined axis dimensions and strides */
22362236
i = idim + combineoffset;
2237-
if (combineoffset < 0 &&
2237+
if (combineoffset < 0 && new_dims[i] != 0 &&
22382238
new_dims[i] != PyArray_DIM(op, idim)) {
22392239
PyErr_Format(PyExc_ValueError,
22402240
"dimensions in operand %d for collapsing "

numpy/core/tests/test_einsum.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -568,6 +568,16 @@ def test_einsum_fixedstridebug(self):
568568
tp = np.tensordot(A,B, axes=(0,0))
569569
assert_equal(es, tp)
570570

571+
def test_einsum_fixed_collapsingbug(self):
572+
# Issue #5147.
573+
# The bug only occured when output argument of einssum was used.
574+
x = np.random.normal(0, 1, (5, 5, 5, 5))
575+
y1 = np.zeros((5, 5))
576+
np.einsum('aabb->ab', x, out=y1)
577+
idx = np.arange(5)
578+
y2 = x[idx[:, None], idx[:, None], idx, idx]
579+
assert_equal(y1, y2)
580+
571581

572582
if __name__ == "__main__":
573583
run_module_suite()

0 commit comments

Comments
 (0)
0