python - numpy:在两个二维数组的一个公共(public)轴上广播乘法

Question

我正在寻找一种方法来分别将两个形状为 (a, b) 和 (b, c) 的二维数组按元素相乘。在两个阵列共有的“b”轴上。

例如，我想要广播(矢量化)的示例是:

import numpy as np    

# some dummy data
A = np.empty((2, 3))
B = np.empty((3, 4))

# naive implementation
C = np.vstack(np.kron(A[:, i], B[i, :]) for i in [0, 1, 2])

# this should give (3, 2, 4)
C.shape

有谁知道在这里做什么？有没有更好的办法？

Answer 1

不同的测试用例:

In [56]: A=np.arange(6).reshape((2,3))
In [57]: B=np.arange(12).reshape((3,4))
In [58]: np.vstack([np.kron(A[:,i],B[i,:]) for i in range(3)])
Out[58]: 
array([[ 0,  0,  0,  0,  0,  3,  6,  9],
       [ 4,  5,  6,  7, 16, 20, 24, 28],
       [16, 18, 20, 22, 40, 45, 50, 55]])

第一次尝试使用 `einsum，保留所有 3 个轴(不求和)

In [60]: np.einsum('ij,jk->ijk',A,B)
Out[60]: 
array([[[ 0,  0,  0,  0],
        [ 4,  5,  6,  7],
        [16, 18, 20, 22]],

       [[ 0,  3,  6,  9],
        [16, 20, 24, 28],
        [40, 45, 50, 55]]])

相同的数字，但形状不同。

我可以对输出轴重新排序，制作一个 2x4x3，它可以 reshape 为 8,3 并转置。

In [64]: np.einsum('ij,jk->ikj',A,B).reshape(8,3).T
Out[64]: 
array([[ 0,  0,  0,  0,  0,  3,  6,  9],
       [ 4,  5,  6,  7, 16, 20, 24, 28],
       [16, 18, 20, 22, 40, 45, 50, 55]])

因此，通过另一次迭代，我可以摆脱转置

In [68]: np.einsum('ij,jk->jik',A,B).reshape(3,8)
Out[68]: 
array([[ 0,  0,  0,  0,  0,  3,  6,  9],
       [ 4,  5,  6,  7, 16, 20, 24, 28],
       [16, 18, 20, 22, 40, 45, 50, 55]])

我应该马上赶到那里。 A 是 (2,3)，B 是 (3,4)，我希望 (3,2,4) reshape 为 (3,8)。 i=2, j=3, k=4 => jik.

所以用另一种方式来描述这个问题，

a_ij * b_jk = c_jik

而且由于我没有使用 einsum 的 sum 部分，常规广播乘法也可以使用，带有一个或多个转置。