Внедрение numpy суммирования по Эйнштейну для суммы точечных произведений

Question

Мне нужно оптимизировать алгоритм, который должен быть максимально быстрым, на данный момент его базовая сумма составляет c для задачи с точечными произведениями 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 раз подряд, поэтому я нашел другое решение barbari c:

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 раз, занимает около секунды, тогда как второй код занимает 0,4 секунды. Тем не менее, я буду использовать их в алгоритме ИИ оптимизации, поэтому он будет повторяться в течение нескольких недель, в этой форме. Переменные формы:

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

Любая помощь, даже без обозначений Эйнштейна, но сокращение 2-го суммирования кода до 4 строк, как указано выше, очень мне поможет.