在沒有明確複製的情況下使用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])