Написать двойную (тройную) сумму как внутренний продукт?

Так как мойnp.dot Ускоряется OpenBlas и Openmpi Интересно, была ли возможность написать двойную сумму

for i in range(N):
     for j in range(N):
         B[k,l] += A[i,j,k,l] * X[i,j]

как внутренний продукт. Прямо сейчас я использую

B = np.einsum("ijkl,ij->kl",A,X)

но, к сожалению, он довольно медленный и использует только один процессор. Есть идеи?

Изменить: я сравнил ответы, данные до сих пор на простом примере, похоже, что они все в одном порядке:

A = np.random.random([200,200,100,100])
X = np.random.random([200,200])
def B1():
    return es("ijkl,ij->kl",A,X) 
def B2():
    return np.tensordot(A, X, [[0,1], [0, 1]])
def B3():
    shp = A.shape
    return np.dot(X.ravel(),A.reshape(shp[0]*shp[1],1)).reshape(shp[2],shp[3])

%timeit B1()
%timeit B2()
%timeit B3()

1 loops, best of 3: 300 ms per loop
10 loops, best of 3: 149 ms per loop
10 loops, best of 3: 150 ms per loop

Исходя из этих результатов, я бы выбрал np.einsum, так как его синтаксис по-прежнему наиболее читабелен, а улучшение по сравнению с двумя другими только в 2 раза. Я предполагаю, что следующим шагом будет вывод кода на С или фортран.

Ответы на вопрос(3)

Ваш ответ на вопрос