Effiziente Reduktion mehrerer Tensoren in Python

Ich habe vier multidimensionale Tensoren v[i,j,k] , a[i,s,l] , w[j,s,t,m] , x[k,t,n] in Numpy, und ich versuche den Tensor z[l,m,n] zu berechnen, der gegeben ist durch:

z[l,m,n] = sum_{i,j,k,s,t} v[i,j,k] * a[i,s,l] * w[j,s,t,m] * x[k,t,n]

Alle Tensoren sind relativ klein (sagen wir weniger als 32k Elemente insgesamt), aber ich muss diese Berechnung viele Male durchführen, deshalb möchte ich, dass die Funktion so wenig Aufwand wie möglich hat.

Ich habe versucht, es mit numpy.einsum wie folgt zu implementieren:

aber es war sehr langsam. Ich habe auch die folgende Sequenz von numpy.tensordot -Aufrufen versucht:

innerhalb einer doppelten for-Schleife, um über s und t zu summieren (sowohl s als auch t sind sehr klein, das ist also kein allzu großes Problem). Dies funktionierte viel besser, aber es ist immer noch nicht so schnell wie ich außer. Ich denke, das kann an all den Operationen liegen, die tensordot intern durchführen muss, bevor das tatsächliche Produkt genommen wird (z. B. das Permutieren der Achsen).

Ich habe mich gefragt, ob es einen effizienteren Weg gibt, diese Art von Operationen in Numpy zu implementieren. Ich hätte auch nichts dagegen, diesen Teil in Cython zu implementieren, aber ich bin mir nicht sicher, welcher Algorithmus der richtige wäre.