summaryrefslogtreecommitdiff
path: root/numpy/core/tests
diff options
context:
space:
mode:
authorCharles Harris <charlesr.harris@gmail.com>2018-02-05 10:33:07 -0700
committerGitHub <noreply@github.com>2018-02-05 10:33:07 -0700
commita1137318a00966e018c0dd048890e8efdde6f436 (patch)
tree8c2bac5e9ba82b79bffb1ccd49300074d407c3e5 /numpy/core/tests
parent1ba96c6cbe77ae077039217850078aa4be0c8287 (diff)
parent6348345b3885dfbee6286eff79e7feccc21e8f41 (diff)
downloadpython-numpy-a1137318a00966e018c0dd048890e8efdde6f436.tar.gz
python-numpy-a1137318a00966e018c0dd048890e8efdde6f436.tar.bz2
python-numpy-a1137318a00966e018c0dd048890e8efdde6f436.zip
Merge pull request #10352 from larsoner/fix-einsum
BUG: Fix einsum optimize logic for singleton dimensions
Diffstat (limited to 'numpy/core/tests')
-rw-r--r--numpy/core/tests/test_einsum.py19
1 files changed, 19 insertions, 0 deletions
diff --git a/numpy/core/tests/test_einsum.py b/numpy/core/tests/test_einsum.py
index d07256ec2..9bd85fdb9 100644
--- a/numpy/core/tests/test_einsum.py
+++ b/numpy/core/tests/test_einsum.py
@@ -481,6 +481,25 @@ class TestEinSum(object):
r = np.arange(4).reshape(2, 2) + 7
assert_equal(np.einsum('z,mz,zm->', p, q, r), 253)
+ # singleton dimensions broadcast (gh-10343)
+ p = np.ones((10,2))
+ q = np.ones((1,2))
+ assert_array_equal(np.einsum('ij,ij->j', p, q, optimize=True),
+ np.einsum('ij,ij->j', p, q, optimize=False))
+ assert_array_equal(np.einsum('ij,ij->j', p, q, optimize=True),
+ [10.] * 2)
+
+ p = np.ones((1, 5))
+ q = np.ones((5, 5))
+ for optimize in (True, False):
+ assert_array_equal(np.einsum("...ij,...jk->...ik", p, p,
+ optimize=optimize),
+ np.einsum("...ij,...jk->...ik", p, q,
+ optimize=optimize))
+ assert_array_equal(np.einsum("...ij,...jk->...ik", p, q,
+ optimize=optimize),
+ np.full((1, 5), 5))
+
def test_einsum_sums_int8(self):
self.check_einsum_sums('i1')