我正在寻找一种方法来逐元素地将两个形状为 (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
有人知道在这里做什么吗？有没有更好的办法？
【问题讨论】：
标签： python arrays numpy optimization
使用不同的测试用例：
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，可以重新整形为 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) 重新整形为 (3,8)。 i=2, j=3, k=4 => jik.
所以另一种描述问题的方式，
a_ij * b_jk = c_jik
而且由于我没有使用einsum 的sum 部分，因此常规广播乘法也可以使用，带有一个或多个转置。
【讨论】：
np.einsum 也能处理这些情况
感谢@hpaulj 对A 和B 的定义
使用np.outer 和np.stack
A = np.arange(6).reshape((2, 3))
B = np.arange(12).reshape((3, 4))
np.stack([np.outer(A[:, i], B[i, :]) for i in range(A.shape[1])])
[[[ 0 0 0 0]
[ 0 3 6 9]]
[[ 4 5 6 7]
[16 20 24 28]]
[[16 18 20 22]
[40 45 50 55]]]
并以正确的形状获得np.einsum
np.einsum('ij, jk->jik', A, B)
[[[ 0 0 0 0]
[ 0 3 6 9]]
[[ 4 5 6 7]
[16 20 24 28]]
[[16 18 20 22]
[40 45 50 55]]]
广播和transpose
(A[:, None] * B.T).transpose(2, 0, 1)
[[[ 0 0 0 0]
[ 0 3 6 9]]
[[ 4 5 6 7]
[16 20 24 28]]
[[16 18 20 22]
[40 45 50 55]]]
形状是(3, 2, 4)
【讨论】：