I have a long array of 3-by-3 matrices, e.g.,
import numpy as np
A = np.random.rand(25, 3, 3)
and for each of the small matrices, I would like to perform an outer product dot(a, a.T). The list comprehension
import numpy as np
B = np.array([
np.dot(a, a.T) for a in A
])
works, but doesn't perform well. A possible improvement could be to do just one big dot product, but I'm having troubles here setting up A correctly for it.
Any hints?
You can obtain the list of transposed matrices as A.swapaxes(1, 2), and the list of products you want as A @ A.swapaxes(1, 2).
import numpy as np
A = np.random.rand(25, 3, 3)
B = np.array([
np.dot(a, a.T) for a in A
])
C = A @ A.swapaxes(1, 2)
(B==C).all() # => True
The @ operator is just syntactic sugar for np.matmul, about which the documentation says that "if either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly."
Looking closely at the code and trying to get a bigger picture reveals some information there. The first axis needs alignment between A and its transposed version as we are iterating through the first axis for both these versions in the code. Then, there is reduction along the third axis in both A and its transposed version. Now, to maintain the alignment and perform such a dot product reduction, there's one good vectorized option in np.einsum. Thus, to solve our case, the implementation would look like this -
B = np.einsum('ijk,ilk->ijl',A,A)
