如何在 numpy 中推广逐元素矩阵运算

How to generalize elementwise matrix operations in numpy

假设我有两组由 3xn 数组 V 和 W 表示的 n 个向量和一组由 3x3xn 数组 Q 表示的 n 个矩阵。我如何向量化运算以得到

a) n 个向量的集合 Q[:,:,k]*V[:,k] for k in range(n)

b) n 个标量的集合 W[:,k].T*Q[:,:,k]*V[:,k] for k in range(n)

这里的乘法理解为矩阵乘法。无法理解我认为应该使用的 einsum 文档。任何 clarification/solution 将不胜感激。

a) the set of n vectors Q[:,:,k]*V[:,k] for k in range(n)

np.einsum('ijk,jk->ik', Q, v)

应该生成一个 (3,n) 数组。矩阵乘积求和结束j.

b) the set of n scalars W[:,k].T*Q[:,:,k]*V[:,k] for k in range(n)

np.einsum('ik,ijk,jk->k', W, Q, V)

我靠记忆工作,我对你需要什么的最佳猜测。所以我的 'ij' 表达式可能需要调整。但试试这些,让我知道它是如何工作的。


测试

In [180]: V=W=np.arange(3*n).reshape(3,n)

In [181]: Q=np.arange(3*3*n).reshape(3,3,n)

In [182]: np.einsum('ijk,jk->ik',Q,V)
Out[182]: 
array([[ 80, 107, 140, 179],
       [224, 287, 356, 431],
       [368, 467, 572, 683]])

In [183]: np.einsum('ik,ijk,jk',W,Q,V)
Out[183]: 28788    # summation over k

In [184]: np.einsum('ik,ijk,jk->k',W,Q,V)
Out[184]: array([ 3840,  5745,  8136, 11067])

有时将 einsum 分成几个步骤会更快,因为它可以防止迭代 space 变得太大。我不认为这里是这种情况,但这就是它的样子。

In [185]: np.einsum('jk,jk->k',np.einsum('ik,ijk->jk',W,Q),V)
Out[185]: array([ 3840,  5745,  8136, 11067])

并使用 Jaime 的评论:

In [186]: np.einsum('i...,ij...,j...',W,Q,V)
Out[186]: array([ 3840,  5745,  8136, 11067])

In [187]: np.einsum('ij...,j...->i...',Q,V)
Out[187]: 
array([[ 80, 107, 140, 179],
       [224, 287, 356, 431],
       [368, 467, 572, 683]])