Added another einsum operation

This commit is contained in:
Alan Saul 2015-04-30 13:59:00 +01:00
parent 55b77064d4
commit 2dcfabf300
2 changed files with 18 additions and 0 deletions

View file

@ -50,3 +50,10 @@ class LinalgTests(np.testing.TestCase):
pure = np.einsum('ijk,jlk->il', A, B)
quick = GPy.util.linalg.ijk_jlk_to_il(A,B)
np.testing.assert_allclose(pure, quick)
def test_einsum_ijk_ljk_to_ilk(self):
A = np.random.randn(150, 20, 5)
B = np.random.randn(20, 30, 5)
pure = np.einsum('ijk,jlk->il', A, B)
quick = GPy.util.linalg.ijk_jlk_to_il(A,B)
np.testing.assert_allclose(pure, quick)

View file

@ -465,3 +465,14 @@ def ijk_jlk_to_il(A, B):
res = np.zeros((A.shape[0], B.shape[1]))
[np.add(np.dot(A[:,:,k], B[:,:,k]), res, res) for k in range(B.shape[-1])]
return res
def ijk_ljk_to_ilk(A, B):
"""
Faster version of einsum np.einsum('ijk,ljk->ilk', A, B)
I.e A.dot(B.T) for every dimension
"""
res = np.empty((A.shape[0], B.shape[0], A.shape[-1]))
[np.dot(A[:,:,i], B[:,:,i].T, res[i,:,:]) for i in range(A.shape[0])]
res = res.swapaxes(0,2)
return res