矩陣乘法與numpy.einsum

Question

我有以下兩個形狀的數組：

    A = (d,w,l)
    B = (d,q)

我想將它們組合成一個帶有形狀的3d數組：

    C = (q,w,l)

更具體一點，在我的情況下，d（3d數組的深度）是2，並且我首先想要將A的上層中的所有位置乘以w * l（所以d = 0）最高行中B的第一個值（所以d = 0，q = 0）。 對於d = 1，我也這樣做，然后將兩者相加：

    C_{q=0,w,l} = A_{d=0,w,l}*B_{d=0,q=0} + A_{d=1,w,l}*B_{d=1,q=0}

我想通過使用numpy.einsum來計算C. 我想到了以下代碼：

    A = np.arange(100).reshape(2,10,5)

    B = np.arange(18).reshape(2,9)

    C = np.einsum('ijk,i -> mjk',A,B)

其中ijk指2,10,5而mjk指9,10,5。 但是我收到了一個錯誤。 有沒有辦法用numpy einsum進行這種乘法？

謝謝

Answer 1

你的形狀A = (d,w,l), B = (d,q), C = (q,w,l)實際上寫出了einsum表達式

C=np.einsum('dwl,dq->qwl',A,B)

我可以測試一下

In [457]: np.allclose(A[0,:,:]*B[0,0]+A[1,:,:]*B[1,0],C[0,:,:])
Out[457]: True