廣播沿B A的每一行，從而避免`repeat`

Question

我有2個二維數組，用1軸相同的尺寸的：

a = np.array(np.arange(6).reshape((2,3))) 
b = np.array(np.arange(12).reshape((3,4))) 

欲相乘，並用b廣播的a每一行，該是

b_r = np.repeat(b[:,:,None], 2, axis=2) 
ab = a.T[:,None,:] * b_r 

是否可以在避開repeat的情況下進行廣播？這個想法是爲了避免repeat操作中不必要的內存分配。

Answer 1

你可以在b[:,:,None]沒有重複，因爲broadcasting它的定義會爲你播放。

因此，簡單地做 -

ab = a.T[:,None,:]*b[:,:,None] 

我們可以使它有點緊湊，雖然跳過了a後:並使用...更換:,:爲b，像這樣 -

ab = a.T[:,None]*b[...,None] 

---

對於踢，這裏是一個使用np.einsum，這會少一點性能，但是一旦我們超過其字符串符號，就會更具表現力 -

ab = np.einsum('ij,jk->jki',a,b)