numpy einsum 折叠除第一维之外的所有内容

问题描述

我有一个维度为 (5,...) 的数组（... 可以是任何东西），我想在 5 之后形成所有维度的点积，以便生成的数组具有形状(5,)。我以为是einsum

i...,i...->i

看起来很有希望，可惜

import numpy as np

a = np.random.rand(5)
b = np.einsum("i...,i...->i",a,a)
assert b.shape == (5,)  # okay


a = np.random.rand(5,2)
b = np.einsum("i...,a)  # fail
assert b.shape == (5,)


a = np.random.rand(5,2,3)
b = np.einsum("i...,)

in einsum
    return c_einsum(*operands,**kwargs)
ValueError: output has more dimensions than subscripts given in einstein sum,but no '...' ellipsis provided to broadcast the extra dimensions.

我或许可以将 a 改造成

b = np.einsum("ij,ij->i",a.reshape(a.shape[0],-1),-1))

但这看起来太乱了。有什么更好的建议吗？

解决方法

根据文档，省略号“启用和控制广播”。您想对这些维度进行点积求和。这些维度需要明确的整体。因此，对于减少最后 2 个维度的 4d，'...ij->...'

In [15]: arr = np.arange(24).reshape(2,3,4)

最后或最后两个维度的减少：

In [17]: np.einsum('...i->...',arr).shape
Out[17]: (2,3)
In [18]: np.einsum('...ij->...',arr).shape
Out[18]: (2,)

sum 等价物：

In [20]: arr.sum(axis=(-1))
Out[20]: 
array([[ 6,22,38],[54,70,86]])
In [21]: arr.sum(axis=(-1)).shape
Out[21]: (2,3)
In [22]: arr.sum(axis=(-2,-1)).shape
Out[22]: (2,)
In [23]: arr.sum(axis=(-2,-1))
Out[23]: array([ 66,210])
In [24]: np.einsum('...ij->...',arr)
Out[24]: array([ 66,210])

来自sympy的这个表达式可能有助于形象化einsum

In [37]: f = Function('f')
    ...: Sum(f(x,y,z),(y,3),(z,2))
Out[37]: 
  2     3             
 ___   ___            
 ╲     ╲              
  ╲     ╲             
  ╱     ╱   f(x,z)
 ╱     ╱              
 ‾‾‾   ‾‾‾            
z = 0 y = 0 

f 是 3 个变量 (3d) 的函数。 einsum 表示法模拟爱因斯坦表示法，其中明确列出了应该求和的维度。那些没有总结的，从某种意义上说，只是随身携带，不需要（如）明确。

,

为了说明，假设张量 A 是一个 2×2 矩阵。你需要的点积可以写成

C_i = A_{kj} A_{ij} D_{ki},

其中 D 是单位矩阵。

使用np.einsum，可以写成

D= np.identity(2)
X= np.einsum("k...,i...,ki->i...",A,D)
C= np.sum(X.reshape(X.shape[0],-1),axis=-1)

我们需要上面的 np.sum 对所有 j 个索引求和。