在numpy.einsum中对省略号广播维度求和

Question

在numpy中,我有一个可以是2-D或3-D的数组,我想将它减少到2-D,同时对每个元素求平方.所以我尝试了这个并不起作用:

A = np.random.rand(5, 3, 3)
np.einsum('...ij,...ij->ij', A, A)

它返回此错误:

ValueError: output has more dimensions than subscripts given in einstein sum, but no '...' ellipsis provided to broadcast the extra dimensions.

我想einsum并不认为当省略号在右侧消失时,我想要省略省略号(如果它们存在).是否有一些"优雅"的方式(即没有检查维度的数量和使用if语句)告诉它我想为3-D做这个:

A = np.random.rand(5, 3, 3)
np.einsum('aij,aij->ij', A, A)

这对于2-D？

A = np.random.rand(3, 3)
np.einsum('ij,ij->ij', A, A)

Answer 1

有时，处理可变维度的“优雅”方法是使用一组if测试，并将它们隐藏在函数调用中。例如，查看np.atleast_3d; 它有一个 4way if/else 子句。我在这里推荐它，只不过它在末尾而不是开头添加了额外的维度。 ifusing 子句reshape并不昂贵（时间方面），所以不要害怕使用它们。即使你发现了一些神奇的功能，也看看它的代码；你可能会对隐藏的内容感到惊讶。

省略号用于“随行”的尺寸，而不是您想要特定控制的尺寸。这里你想要对初始维度求和，所以你需要显式地对其进行索引：

In [161]: np.einsum('i...,i...',A,A)
Out[161]: 
array([[ 1.26942035,  1.32052776,  1.74118617],
       [ 1.59679765,  1.49331565,  2.04573002],
       [ 2.29027005,  1.48351522,  1.36679208]])
In [162]: np.einsum('aij,aij->ij',A,A)
Out[162]: 
array([[ 1.26942035,  1.32052776,  1.74118617],
       [ 1.59679765,  1.49331565,  2.04573002],
       [ 2.29027005,  1.48351522,  1.36679208]])

对于二维数组：

In [165]: np.einsum('ij,ij->ij',A[0],A[0])
Out[165]: 
array([[ 0.20497776,  0.11632197,  0.65396968],
       [ 0.0529767 ,  0.24723351,  0.27559647],
       [ 0.62806525,  0.33081124,  0.57070406]])
In [166]: A[0]*A[0]
Out[166]: 
array([[ 0.20497776,  0.11632197,  0.65396968],
       [ 0.0529767 ,  0.24723351,  0.27559647],
       [ 0.62806525,  0.33081124,  0.57070406]])
In [167]: 
In [167]: np.einsum('...,...',A[0],A[0])
Out[167]: 
array([[ 0.20497776,  0.11632197,  0.65396968],
       [ 0.0529767 ,  0.24723351,  0.27559647],
       [ 0.62806525,  0.33081124,  0.57070406]])

我认为你不能用一种表达来处理这两种情况。

获得第一笔款项的另一种方法

In [168]: (A*A).sum(axis=0)
Out[168]: 
array([[ 1.26942035,  1.32052776,  1.74118617],
       [ 1.59679765,  1.49331565,  2.04573002],
       [ 2.29027005,  1.48351522,  1.36679208]])

我贡献了修复省略号处理的补丁，但那是几年前的事了。所以这些细节在我的脑海里已经不太新鲜了。作为其中的一部分，我对字符串表达式的解析进行了逆向工程（原始代码已编译），并且如果我们需要更明确的答案，可以查看该代码（或向您推荐它）。

In [172]: np.einsum('...ij,...ij->ij',A,A)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-172-dfe39e268402> in <module>()
----> 1 np.einsum('...ij,...ij->ij',A,A)

ValueError: output has more dimensions than subscripts given in 
 einstein sum, but no '...' ellipsis provided to broadcast 
 the extra dimensions.
In [173]: np.einsum('...ij,...ij->...ij',A,A).shape
Out[173]: (5, 3, 3)

错误消息表明它正在尝试将...尺寸传递到输出，但不能 - 因为输出缺少尺寸或.... 换句话说，它不对...维度执行求和。它们不变地传递到输出（适用广播规则）。