diff options
author | Charles Harris <charlesr.harris@gmail.com> | 2018-02-05 10:33:07 -0700 |
---|---|---|
committer | GitHub <noreply@github.com> | 2018-02-05 10:33:07 -0700 |
commit | a1137318a00966e018c0dd048890e8efdde6f436 (patch) | |
tree | 8c2bac5e9ba82b79bffb1ccd49300074d407c3e5 /numpy/core/tests | |
parent | 1ba96c6cbe77ae077039217850078aa4be0c8287 (diff) | |
parent | 6348345b3885dfbee6286eff79e7feccc21e8f41 (diff) | |
download | python-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.py | 19 |
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') |