Numpy - 两个矩阵的逐行外积

Question

我有两个numpy数组: 形状为(b, i)的A和形状为(b, o)的B。我想计算一个形状为(b, i, o)的数组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数组。

有没有更有效的方法？

- Marvin Lerousseau

---

请参阅如何创建一个最小、完整和可验证的示例。 - keepAlive

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)

- MB-F

---

谢谢，非常完美。我会尽快接受您的答案。 - Marvin Lerousseau

有没有办法让这个适用于稀疏矩阵？ - Thomas Ahle

@ThomasAhle 我自己从未尝试过，但是使用稀疏后端的 opt_einsum 看起来很有前途。 - MB-F

Answer 2

为什么不简单地

A[:, :, None] * B[:, None, :]

根据您的惯例和数据类型，您可能需要在其他地方添加np.conj。请注意，np.newaxis只是None。

- dseuss