diff --git a/GPy/testing/linalg_test.py b/GPy/testing/linalg_test.py index 81cb0368..ad524156 100644 --- a/GPy/testing/linalg_test.py +++ b/GPy/testing/linalg_test.py @@ -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) diff --git a/GPy/util/linalg.py b/GPy/util/linalg.py index 285701e7..bc047654 100644 --- a/GPy/util/linalg.py +++ b/GPy/util/linalg.py @@ -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