python - 用于多个外积的 Numpy 运算符

Question

import numpy as np
mat1 = np.random.rand(2,3)
mat2 = np.random.rand(2,5)

我希望得到一个 2x3x5 张量，其中每一层都是通过将 mat1 的 3x1 转置行乘以 mat2 的 1x5 行获得的 3x5 外积。

可以用 numpy matmul 来完成吗？

Answer 1

您可以简单地使用 broadcasting使用 np.newaxis/None 扩展其尺寸后-

mat1[...,None]*mat2[:,None]

这将是性能最高的，因为这里不需要 sum-reduction 来保证来自 np.einsum 或 np.matmul 的服务.

如果你还想拖入np.matmul ，它与广播基本相同:

np.matmul(mat1[...,None],mat2[:,None])

与 np.einsum ，如果您熟悉它的字符串表示法，它可能看起来比其他的更整洁 -

np.einsum('ij,ik->ijk',mat1,mat2)
#          23,25->235  (to explain einsum's string notation using axes lens)