python - 高效计算多个张量内积

Question

我正在使用一个 k x k x k x k 张量(比如 S)和一个大小为 (n, k) 的数组 X 。粗略地说，X 的行对应于图的节点特征。对于每对边(例如 e = (u, v) 和 e' = (u_, v_))，我想计算一个新元素，如下所示:

elt = np.sum(S * np.multiply.outer(np.outer(X[u, :], X[v, :]), np.outer(X[u_, :], X[v_, :])))

我想知道是否有一种方法可以更有效地完成此操作，而不是在索引上使用 4 个嵌套循环。如果我只处理节点对，并且 S 只是一个 k x k 矩阵，则可以简单地写为

all_elts = X @ S @ X.T

但是，我不确定这如何在多个维度上推广。非常感谢任何帮助!

Answer 1

这是一个展示如何使用einsum()的示例:

import numpy as np
from itertools import product

n = 4
x = np.random.randn(n, n)
S = np.random.randn(n, n, n, n)

res = np.zeros((n, n, n, n))
for i, j, k, l in product(range(n), range(n), range(n), range(n)):
    res[i, j, k, l] = np.sum(S * np.multiply.outer(np.outer(x[i, :], x[j, :]), np.outer(x[k, :], x[l, :])))

res2 = np.einsum("efgh,ae,bf,cg,dh->abcd", S, x, x, x, x)
np.allclose(res, res2)