理解两个张量的np.einsum('ijk,jil->kl')元素求和过程细节

我们先定义两个数值张量：

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

目标是实现爱因斯坦求和：

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']]]

拆解求和逻辑 #

爱因斯坦求和的核心是下标对应规则：

• 张量a的下标ijk：i对应a的第0维度（长度4），j对应第1维度（长度2），k对应第2维度（长度1）
• 张量b的下标jil：j对应b的第0维度（长度2），i对应b的第1维度（长度4），l对应b的第2维度（长度2）
• 输出的kl表示结果只保留k和l维度，重复出现的下标i和j是求和维度（对所有可能的i、j取值做累加）

简单来说，输出张量的每个元素c[k][l]，等于所有i（0-3）、j（0-1）对应的a[i][j][k] * b[j][i][l]的总和。

用字符串张量演示具体计算 #

以输出的两个元素为例：

• c[0][0]的计算（k=0，l=0）：

  e*x + r*s + t*v + y*a + u*m + o*j + p*f + q*l
  

  对应逐个i、j的组合：

  • i=0,j=0: aa[0][0][0] * bb[0][0][0] = e*x
  • i=0,j=1: aa[0][1][0] * bb[1][0][0] = r*s
  • i=1,j=0: aa[1][0][0] * bb[0][1][0] = t*v
  • i=1,j=1: aa[1][1][0] * bb[1][1][0] = y*a
  • i=2,j=0: aa[2][0][0] * bb[0][2][0] = u*m
  • i=2,j=1: aa[2][1][0] * bb[1][2][0] = o*j
  • i=3,j=0: aa[3][0][0] * bb[0][3][0] = p*f
  • i=3,j=1: aa[3][1][0] * bb[1][3][0] = q*l

• c[0][1]的计算（k=0，l=1）：

  e*c + r*w + t*n + y*z + u*h + o*k + p*d + q*b
  

  对应逐个i、j的组合：

  • i=0,j=0: aa[0][0][0] * bb[0][0][1] = e*c
  • i=0,j=1: aa[0][1][0] * bb[1][0][1] = r*w
  • i=1,j=0: aa[1][0][0] * bb[0][1][1] = t*n
  • i=1,j=1: aa[1][1][0] * bb[1][1][1] = y*z
  • i=2,j=0: aa[2][0][0] * bb[0][2][1] = u*h
  • i=2,j=1: aa[2][1][0] * bb[1][2][1] = o*k
  • i=3,j=0: aa[3][0][0] * bb[0][3][1] = p*d
  • i=3,j=1: aa[3][1][0] * bb[1][3][1] = q*b

这样就能清晰看到，np.einsum('ijk,jil->kl')不是单纯转置b后运算，而是按照下标匹配规则，对所有重叠维度的元素乘积做累加，最终得到k×l维度的结果。

内容的提问来源于stack exchange，提问作者unbelievable