Python 实现点积和的numpy-einstein求和

我需要优化一个算法，它需要尽可能快，现在，它是一个基本的2向量点积和问题，但我认为我的解决方案有点多余，爱因斯坦表示法可以得到更快的结果。首先，我有一个单一的价值案例：

for t in range(2):
  for i2 in range(40):
     for j2 in range(40):
       U_out[i2,j2,t] += np.sum(np.multiply(U_in[:,:,t],constants[:,:,j2,i2,t]))

我需要连续使用8次，所以我想出了另一个野蛮的解决方案：

for t in range(2):
  for i2 in range(40):
    for j2 in range(40):
      U_out[i2,j2,t,0] += np.sum(np.multiply(U_in[:,:,t,0],constants[:,:,j2,i2,t]))
      U_out[i2,j2,t,1] += np.sum(np.multiply(U_in[:,:,t,1],constants[:,:,j2,i2,t]))
      U_out[i2,j2,t,2] += np.sum(np.multiply(U_in[:,:,t,2],constants[:,:,j2,i2,t]))
      U_out[i2,j2,t,3] += np.sum(np.multiply(U_in[:,:,t,3],constants[:,:,j2,i2,t]))
      U_out[i2,j2,t,4] += np.sum(np.multiply(U_in[:,:,t,4],constants[:,:,j2,i2,t]))
      U_out[i2,j2,t,5] += np.sum(np.multiply(U_in[:,:,t,5],constants[:,:,j2,i2,t]))
      U_out[i2,j2,t,6] += np.sum(np.multiply(U_in[:,:,t,6],constants[:,:,j2,i2,t]))
      U_out[i2,j2,t,7] += np.sum(np.multiply(U_in[:,:,t,7],constants[:,:,j2,i2,t]))

现在，第一个代码重复了8次，大约是1秒，而第二个代码需要0.4秒。然而，我将在优化AI算法中使用它们，因此它将以这种形式循环数周。可变形状包括：

U_out = (40,40,2) 1st code , (40,40,2,8) 2nd code 
U_in  = (40,40,2) 1st code , (40,40,2,8) 2nd code 
constants = (40,40,40,40,2) for both codes

---

任何帮助，即使没有爱因斯坦符号，但将第二个代码总和减少到4行（如上所述）会对我有很大帮助。

如果你能用唯一的值或符号显示数组的大小，einsum可能会自己写。你说的“可能自己写”是什么意思看起来像

U in

有形状（M，N，T，Q）<代码>常数（M，N，K，L，T）和

U out

（K，L，T，Q），您正在对M，N进行

dot

求和。如果是这样，einsum索引表达式将是'mntq，mnklt->kltq'。非常感谢，这非常快！你可能想看看。