了解两个张量的 einsum 应用细节

Question

我们有两个张量：

 a = np.arange(8.).reshape(4,2,1)
 b = np.arange(16.).reshape(2,4,2)

我们要严正

np.einsum('ijk,jil->kl', a, b)

虽然我们可以得到它的结果，但我们仍然坚持了解张量元素求和的过程细节。

首先我们知道如何

np.einsum('jil', b)

更改

b

张量的元素顺序。

但是我们无法理解

np.einsum('ijk,jil->kl', a, b)

如何组合（求和）张量元素。

为了跟踪进程，我们使用了字符串：

aa=[[['e'],
  ['r']],

 [['t'],
  ['y']],

 [['u'],
  ['o']],

 [['p'],
  ['q']]]

和

bb=[[[ 'x', 'c'],
  [ 'v' , 'n'],
  [ 'm',  'h'],
  [ 'f' , 'd']],

 [[ 's',  'w'],
  [ 'a','z'],
  ['j', 'k'],
  ['l', 'b']]]

因为我们想看看不同的元素如何组合以获得

np.einsum('ijk,jil->kl', aa, bb)

！

但是

np.einsum('jil', bb)

工作正常，但它没有向我显示元素求和的细节。

Answer 1

有几种方法可以理解这一点。

一种是使用 @Onyambu 建议的示例。

>>> np.einsum('ijk,jil->ijkl', a, b)

array([[[[  0.,   0.]],
        [[  8.,   9.]]],

       [[[  4.,   6.]],
        [[ 30.,  33.]]],

       [[[ 16.,  20.]],
        [[ 60.,  65.]]],

       [[[ 36.,  42.]],
        [[ 98., 105.]]]])

通过将 i 和 j 作为输出的索引，输出数组的形状不再是 (k, l)，而是 (i, j, k, l)。此外，没有任何相乘的元素被加在一起。输出数组的每个元素都是原始数组的两个元素之和。

要回到原始行为，我们可以按轴 1 求和：

>>> np.einsum('ijk,jil->ijkl', a, b).sum(axis=1)
array([[[  8.,   9.]],
       [[ 34.,  39.]],
       [[ 76.,  85.]],
       [[134., 147.]]])

然后按轴0求和：

>>> np.einsum('ijk,jil->ijkl', a, b).sum(axis=1).sum(axis=0)

array([[252., 280.]])

理解这一点的另一种方法是将其转换为显式循环。

以下代码与此 einsum 等效，但速度较慢。 （它也不检查 A 和 B 的形状是否兼容。）

def sum_array(A, B):
    i_len, j_len, k_len = A.shape
    _, _, l_len = B.shape
    
    ret = np.zeros((k_len, l_len))
    for i in range(i_len):
        for j in range(j_len):
            for k in range(k_len):
                for l in range(l_len):
                    ret[k, l] += A[i, j, k] * B[j, i, l]
    return ret

这给了我们相同的结果，

array([[252., 280.]])

。

注意循环的内行

ret[k, l] += A[i, j, k] * B[j, i, l]

与 einsum 下标

'ijk,jil->kl'

相似，不同之处在于

kl

已移至开头，并且

ijk

用于索引 A，并且

jil 

正在用于索引 B。

更多信息

理解 NumPy 的 einsum