python - 在 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)

这是二维的吗？

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

Answer 1

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

---

省略号用于“顺其自然”的维度，而不是您想要特定控制的维度。这里你想对初始维度求和，所以你需要明确地索引它:

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)

错误消息说它正在尝试将 ... 维度传递给输出，但不能 - 因为输出缺少维度或 ...。换句话说，它不会对 ... 维度执行求和。它们传递给输出不变(适用广播规则)。