Numpy矩阵乘法广播

Question

我有一个由N 3x3数组组成的数组（矩阵的集合，虽然数据类型是np.ndarray），我有一个由N 3x1数组（向量集合）组成的数组。 我想要做的是将每个矩阵乘以每个向量，所以我希望得到N个3x1数组。

简单的例子：

A = np.ones((6,3,3))
B = np.ones((6,3,1))
np.dot(A,B) # This gives me a 6x3x6x1 array, which is not what I want
np.array(map(np.dot,A,B)) # This gives me exactly what I want, but I don't want to have to rely on map

我已经厌倦了各种各样的重塑，探索了einsum等等，但是不能按照我想要的方式让它工作。 如何使用numpy广播？ 这个操作最终需要完成数千次，而且我不希望map或列表理解操作减慢速度。

Answer 1

您可以使用np.einsum计算点积并创建所需形状的矩阵：

np.einsum('ijk,ikl->ijl', A, B)

Answer 2

A = np.random.rand(6, 3, 3)
B = np.random.rand(6, 3, 1)
C = np.array(map(np.dot, A, B))
D = np.sum(A*B.swapaxes(1, 2), axis=2)[..., None]
assert np.allclose(C, D)
assert C.shape == D.shape == (6, 3, 1)

“allclose”是因为在1e-16的数量级上，两种方法之间存在一些浮点舍入差异。

.swapaxis和[...，None]只是为了让数组符合你指定的形状。 您还可以使用以下代码更简单地表示它：

A = np.random.rand(6, 3, 3)
B = np.random.rand(6, 3)
C = np.array(map(np.dot, A, B))
D = np.sum(A*B[:, None, :], axis=2)
assert np.allclose(C, D)
assert C.shape == D.shape == (6, 3)

Answer 3

可以使用PEP 465中引入的Python 3.5或更高版本中的内置矩阵乘法。

$ python --version
Python 3.6.6

>>> import numpy as np
>>> A = np.ones((6,3,3))
>>> B = np.ones((6,3,1))
>>> C = A @ B
>>> print(C)
[[[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]]