矩阵乘法与numpy.einsum

时间:2015-05-01 14:45:32

标签: python arrays numpy

我有以下两个形状的数组:

    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进行这种乘法?

由于

1 个答案:

答案 0 :(得分:4)

你的形状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