I have two multidimensional NumPy arrays, A
and B
, with A.shape = (K, d, N)
and B.shape = (K, N, d)
. I would like to perform an element-wise operation over axis 0 (K
), with that operation being matrix multiplication over axes 1 and 2 (d, N
and N, d
). So the result should be a multidimensional array C
with C.shape = (K, d, d)
, so that C[k] = np.dot(A[k], B[k])
. A naive implementation would look like this:
C = np.vstack([np.dot(A[k], B[k])[np.newaxis, :, :] for k in xrange(K)])
but this implementation is slow. A slightly faster approach looks like this:
C = np.dot(A, B)[:, :, 0, :]
which uses the default behaviour of np.dot
on multidimensional arrays, giving me an array with shape (K, d, K, d)
. However, this approach computes the required answer K
times (each of the entries along axis 2 are the same). Asymptotically it will be slower than the first approach, but the overhead is much less. I am also aware of the following approach:
from numpy.core.umath_tests import matrix_multiply
C = matrix_multiply(A, B)
but I am not guaranteed that this function will be available. My question is thus, does NumPy provide a standard way of doing this efficiently? An answer which applies to multidimensional arrays in general would be perfect, but an answer specific to only this case would be great too.
Edit: As pointed out by @Juh_, the second approach is incorrect. The correct version is:
C = np.dot(A, B).diagonal(axis1=0, axis2=2).transpose(2, 0, 1)
but the overhead added makes it slower than the first approach, even for small matrices. The last approach is winning by a long shot on all my timing tests, for small and large matrices. I'm now strongly considering using this if no better solution crops up, even if that would mean copying the numpy.core.umath_tests
library (written in C) into my project.