在没有明确复制的情况下使用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])