NumPy: produit Dot pour de nombreuses petites matrices à la fois

Question

J'ai une longue série de 3-en-3 matrices, par exemple,

import numpy as np 

A = np.array([ 
    [[1, 2, 3], 
    [3, 4, 5], 
    [4, 5, 1]], 
    [[2, 0, 3], 
    [5, 1, 5], 
    [4, 2, 1]] 
    ]) 

et pour chacune des petites matrices, je voudrais effectuer un produit externe dot(a, a.T). La compréhension de la liste

import numpy as np 

B = np.array([ 
    np.dot(a, a.T) for a in A 
    ]) 

fonctionne, mais ne fonctionne pas bien. Une amélioration possible pourrait être de faire un gros dot produit, mais j'ai des problèmes ici à mettre en place A correctement pour cela.

Des indices?

Answer 1

Vous pouvez obtenir la liste des matrices transposés comme A.swapaxes(1, 2), et la liste des produits que vous voulez que A @ A.swapaxes(1, 2).

import numpy as np 

A = np.array([ 
    [[1, 2, 3], 
    [3, 4, 5], 
    [4, 5, 1]], 
    [[2, 0, 3], 
    [5, 1, 5], 
    [4, 2, 1]] 
    ]) 

B = np.array([ 
    np.dot(a, a.T) for a in A 
    ]) 

C = A @ A.swapaxes(1, 2) 

(B==C).all()  # => True 

Le @ operator est le sucre juste syntaxique pour np.matmul, dont le documentation dit que «si l'argument est ND, N> 2, elle est traitée comme une pile de matrices résidant dans les deux derniers indices et diffusion en conséquence "

Answer 2

En regardant de près le code et en essayant d'obtenir une plus grande image révèle des informations là-bas. Le premier axe a besoin d'un alignement entre A et sa version transposée car nous parcourons le premier axe pour ces deux versions dans le code. Ensuite, il y a une réduction le long du troisième axe à la fois dans A et sa version transposée. Maintenant, pour maintenir l'alignement et effectuer une telle réduction de produit scalaire, il y a une bonne option vectorisée dans np.einsum. Ainsi, pour résoudre notre cas, la mise en œuvre ressemblerait à ceci -

B = np.einsum('ijk,ilk->ijl',A,A)