python-2.7 - numpy-两个矩阵的逐行外积

Question

我有两个numpy数组:形状为(b，i)的A和形状为(b，o)的B。我想计算一个形状为(b，i，o)的数组R，其中R的每一行l都包含A行l和B行l的外积。到目前为止，我所拥有的是:

import numpy as np

A = np.ones((10, 2))
B = np.ones((10, 6))
R = np.asarray([np.outer(a, b) for a, b in zip(A, B)])
assert R.shape == (10, 2, 6)

我认为这种方法太慢了，因为zip和最终转换为numpy数组。

有更有效的方法吗？

Answer 1

numpy.matmul 可以实现，它可以对“矩阵堆栈”进行乘法运算。在这种情况下，我们想将一列列向量与一列行向量相乘。首先将矩阵A成形为(b，i，1)，将B成形为(b，1，o)。然后使用matmul进行b乘以外部乘积:

import numpy as np

i, b, o = 3, 4, 5

A = np.ones((b, i))
B = np.ones((b, o))

print(np.matmul(A[:, :, np.newaxis], B[:, np.newaxis, :]).shape)  # (4, 3, 5)

一种替代方法是使用 numpy.einsum ，它可以直接表示您的索引符号:

np.einsum('bi,bo->bio', A, B)