Numpy ndarray Multiplikation

Question

Ich habe zwei 3D

A=np.array([[[1, 1], 
      [1, 1], 
      [1, 1]], 

      [[2, 2], 
      [2, 2], 
      [2, 2]]]) 

B=np.array([[[ 2, 0], 
      [ 0, 2]], 

      [[ 2, -2], 
      [-2, 2]]]) 

ich die AB-Array mit Elementen AB _ijk = Σ _m (A _IJM * B _imk), wo die ndarray numpy erstellen möchten Summierung ist nur über den M-Index (wiederholt) und nicht über i (das wird wiederum wiederholt).

Mit anderen Worten kann ich di AB ndarray mit diesem erhält für Schleife

for i in range(2): 
    AB[i,:,:]=np.dot(A[i,:,:],B[i,:,:]) 

und AB gleich

array([[[ 2., 2.], 
    [ 2., 2.], 
    [ 2., 2.]], 

    [[ 0., 0.], 
    [ 0., 0.], 
    [ 0., 0.]]]) 

Gibt es Art und Weise die for-Schleife zu vermeiden? Wie kann ich das AB-Array mit Tensordot oder Einsum erhalten?

Vielen Dank für Ihre Antworten, ich schätze es wirklich.

Answer 1

ABijk=∑m (Aijm*Bimk) übersetzt

AB = np.einsum('ijm,imk->ijk', A, B) 

ich denke, der matmul Operator wird auch diesen

AB = A @ B 

handhaben, da es die normalen dot auf den in den letzten zwei Dimensionen nimmt, zusammen als Freigepäck, den Rest tragen .

Testen Sie diese und lassen Sie mich wissen, wenn sie funktionieren.

Answer 2

Auf einem ausreichend letzten NumPy (1.10+), können Sie

AB = np.matmul(A, B) 

tun oder (wenn Sie auch Python 3.5 und höher):

AB = A @ B 

Wenn Sie nicht über 1,10 NumPy +, können Sie

AB = np.einsum('ijm,imk->ijk', A, B) 

Für große J/M/K Dimensionen tun, vor allem, wenn Sie ein gutes BLAS haben, kann es auch sinnvoll sein, die expl bedenkt icit for Schleife mit dot. Die BLAS-Matrix multipliziert möglicherweise mehr Zeit als der Overhead von mehr interpretiertem Python verliert. Ich denke, np.matmul und @ sollen die gleichen Dinge nutzen dot tut, aber ich glaube nicht, np.einsum tut.