Numpy ndarray умножение

Question

У меня есть два 3D Numpy ndarray

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]]]) 

Я хочу создать массив AB с элементами AB _IJK = Σ _м (A _IJM * B _ИМК), где суммирование ведется только по m-индексу (повторяется) и не превышает i (что, в свою очередь, повторяется).

Другими словами, я могу получить ди AB ndarray с этим цикл

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

и AB равна

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

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

Есть ли способ избежать для цикла? Как я могу получить массив AB с тензордотом или einsum?

Благодарим за ответы, я очень благодарен вам.

Answer 1

ABijk=∑m (Aijm*Bimk) переводит

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

Я думаю, что оператор matmul обрабатывает эту

AB = A @ B 

, поскольку он занимает нормальное dot на последних 2-х измерениях, неся отдых вместе как бесплатный багаж ,

Проверьте те и сообщите мне, если они будут работать.

Answer 2

На достаточно недавнем NumPy (1.10+), вы можете сделать

AB = np.matmul(A, B) 

или (если у вас также есть Python 3.5+):

AB = A @ B 

Если вы не имеете NumPy 1,10 +, вы можете сделать

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

Для больших J/размеры M/K, особенно если у вас есть хорошие BLAS, это также может быть стоит рассмотреть EXPL icit for цикл с dot. Масштабирование матрицы BLAS может сэкономить больше времени, чем накладные расходы на более интерпретируемые потери Python. Думаю, np.matmul и @ предполагается использовать одни и те же вещи dot делает, но я не думаю, np.einsum делает.