python - 在没有显式复制的情况下使用 einsum

Question

我有一个形状为 (n, n) 的矩阵 A 和另一个形状为 (p, n) 的矩阵 b。我需要得到一个矩阵 C 使得

C[i] = (A * b[i, np.newaxis, :]) * b[i, :, np.newaxis]

我目前正在通过执行 np.matlib.repmat(A) 然后使用 np.einsum 来做到这一点，如下所示

A1 = np.matlib.repmat(A, p, 1).reshape(p, n, n)
C = np.einsum('ijk, ij..., ik... -> ijk', A1, b, b)

但我可以看到创建 A1 是一种资源浪费，因为它只是一遍又一遍地使用相同的值。无论如何我可以在不创建中间矩阵 A1 的情况下做到这一点吗？

Answer 1

np.einsum像这样使用-

np.einsum('ij,kij,kij->kij',A,b[:,None,:],b[...,None])