将矩阵产品沿着最后两个维度取出的单一操作

对于一对由分别a和b表示形状(K, N, M)和(K, M, N)的numpy的阵列，有一种方法来计算下列作为单个矢量化操作：

import numpy as np 
K = 5 
N = 2 
M = 3 
a = np.random.randn(K, N, M) 
b = np.random.randn(K, M, N) 

output = np.empty((K, N, N)) 
for each_a, each_b, each_out in zip(a, b, output): 
    each_out[:] = each_a.dot(each_b)

一个简单a.dot(b)返回点积对于每一对第一个轴（所以它返回一个形状数组(K, N, K, N)。

编辑：充实了代码，让那些无法理解问题的人有点为难

来源

2014-09-29 Henry Gomersall

我不完全清楚你想要达到什么目的，但是你不是在寻找'a.dot（b）'吗？ – Wolph 2014-09-29 18:36:10

不，如上所述，'a.dot（b）'返回形状为'（K，N，K，N）'的数组，即每对第一个轴的点积。我想要达到的是代码。这不清楚吗？ – 2014-09-29 18:40:12

我只有时间快速浏览一下，没有测试的情况下（和我提出的解决方案）之间的差异对我而言并不清楚。但它看起来像farenorth已经有一个很好的解决方案:) – Wolph 2014-09-29 19:50:36

我回答了一个类似的问题，回来后：Element-wise matrix multiplication in NumPy。

我想你要找的是什么：

output = np.einsum('ijk,ikl->ijl', a, b)

祝你好运！

来源

2014-09-29 19:02:38 farenorth

谢谢，这完美的作品。我玩了一点Einsum，但一天结束的时候，脑fra停了一个合适的答案:) – 2014-09-29 20:41:37

这是一个'einsum'省略号可以让你编写一个表达式在2d和3d一样工作的情况：''... jk ... KL - > ... jl'' – hpaulj 2014-09-29 21:50:11

将矩阵产品沿着最后两个维度取出的单一操作

回答

相关问题