高效计算多个张量内积

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)

高效计算多个张量内积

Efficiently computing multiple tensor inner products

python

numpy

scipy

multidimensional-array

tensor