python - Numpy 相当于点(A，B，3)

Question

假设我有两个 3 维矩阵，像这样(取自这个 matlab 示例 http://www.mathworks.com/help/matlab/ref/dot.html ):

A = cat(3,[1 1;1 1],[2 3;4 5],[6 7;8 9])
B = cat(3,[2 2;2 2],[10 11;12 13],[14 15; 16 17])

如果我想沿三维获取成对的点积，我可以在 matlab 中这样做:

C = dot(A,B,3)

这会给出结果:

C =
  106   140
  178   220

numpy 中的等效操作是什么，最好是矢量化选项，以避免必须在整个数组中编写双 for 循环。我似乎无法理解 np.tensordot 或 np.inner 应该做什么，但它们可能是选项。

Answer 1

In [169]:

A = np.dstack([[[1, 1],[1 ,1]],[[2 ,3],[4, 5]],[[6, 7],[8, 9]]])
B = np.dstack([[[2, 2],[2, 2]],[[10, 11],[12, 13]],[[14, 15], [16, 17]]])
c=np.tensordot(A, B.T,1)
np.vstack([np.diag(c[:,i,i]) for i in range(A.shape[0])]).T
Out[169]:
array([[106, 140],
       [178, 220]])

但令人惊讶的是它是最慢的:

In [170]:

%%timeit
c=np.tensordot(A, B.T,1)
np.vstack([np.diag(c[:,i,i]) for i in range(A.shape[0])]).T
10000 loops, best of 3: 95.2 µs per loop
In [171]:

%timeit np.einsum('i...,i...',a,b)
100000 loops, best of 3: 6.93 µs per loop
In [172]:

%timeit inner1d(A,B)
100000 loops, best of 3: 4.51 µs per loop