没有张量点积的向量化 numpy 计算

Question

我想将以下数学公式的特定情况（来自本文的表 2 和附录 A ）矢量化numpy：

我想计算的情况如下，可以忽略平方根下的比例因子。

该术语w_kij - w_ij_bar是一个n x p x p矩阵，其中n通常远大于p。

我实现了两种解决方案，它们都不是特别好：一种涉及双循环，而另一种会很快用不必要的计算填充内存。

dummy_data = np.random.normal(size=(100, 5, 5))

# approach 1: a double loop

out_hack = np.zeros((5, 5))

for i in range(5):

    for j in range(5):

        out_hack[i, j] = (dummy_data.T[j, j, :]*dummy_data[:, j, i]).sum()

# approach 2: slicing a diagonal from a tensor dot product

out = np.tensordot(dummy_data.T, dummy_data, axes=1)

out = out.diagonal(0, 0, 2).diagonal(0, 0, 2)

print((out.round(6) == out_hack.round(6)).all())

>>> True

有没有办法在这两种方法之间找到中间立场？

Answer 1

np.einsum几乎从字面上翻译——np.einsum('kjj,kji->ij',dummy_data,dummy_data)