python中带广播的元素多维矩阵乘法

Question

是否有一种有效的(numpy函数)方法来对两个不同大小的数组进行元素矩阵乘法，并将其广播到一个新数组中。例如,

a = np.arange(24).reshape((2,12)) #gives a 2x12 array 
b = np.arange(36).reshape((3,12)) #gives a 3x12 array

然后沿着维度乘以12(不求和)得到最终的三维矩阵"c"形状为2x3x12，其中

c[0,0,:] = [a[0,0]*b[0,0], a[0,1]*b[0,1], ... a[0,11]*b[0,11]]
c[1,0,:] = [a[1,0]*b[0,0], a[1,1]*b[0,1], ... a[1,11]*b[0,11]]
c[0,1,:] = [a[0,0]*b[1,0], a[0,1]*b[1,1], ... a[0,11]*b[1,11]]

我可以得到我想要的:

a2 = np.repeat(a[..., np.newaxis], 3, axis=2)
b2 = b.T[np.newaxis, ...]
c = np.swapaxes(a2*b2, 1, 2)

输出:

array([[[  0,   1,   4,   9,  16,  25,  36,  49,  64,  81, 100, 121],
[  0,  13,  28,  45,  64,  85, 108, 133, 160, 189, 220, 253],
[  0,  25,  52,  81, 112, 145, 180, 217, 256, 297, 340, 385]],
[[  0,  13,  28,  45,  64,  85, 108, 133, 160, 189, 220, 253],
[144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484, 529],
[288, 325, 364, 405, 448, 493, 540, 589, 640, 693, 748, 805]]])

但是对于一些我认为不应该太少见的东西，使用5个类似numpy的命令感觉真的很低效。

Answer 1

当然，你可以使用广播:

out = a[:, None, :] * b[None, :, :]

:

array([[[  0,   1,   4,   9,  16,  25,  36,  49,  64,  81, 100, 121],
[  0,  13,  28,  45,  64,  85, 108, 133, 160, 189, 220, 253],
[  0,  25,  52,  81, 112, 145, 180, 217, 256, 297, 340, 385]],
[[  0,  13,  28,  45,  64,  85, 108, 133, 160, 189, 220, 253],
[144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484, 529],
[288, 325, 364, 405, 448, 493, 540, 589, 640, 693, 748, 805]]])