vectorized multi-matrix multiplication
I have two lists of 3x3 arrays and I would like to compute the matrix product of the i-th element in the first list with the i-th element in the second list. Of course, I could just loop over the lists:
for i in range(n): out[i] = dot( matrices1[i], matrices2[i] )
However, the list is quite long, and each matrix is very small (3x3) so this turns out to be quite slow. Is there a way to do this with a single numpy call? I have looked at tensordot but it outputs an N x N x 3 x 3 array, whereas I want an N x 3 x 3 output. I've also looked at various broadcasting tricks but I haven't found anything that works yet.
Alex
Assuming matrices1 and matrices2 are actually arrays of size (N, 3, 3) you can do:
np.einsum('nij,njk->nik', matrices1, matrices2)
On Sun, Aug 26, 2012 at 11:04 AM, Alex Flint alex.flint@gmail.com wrote:
I have two lists of 3x3 arrays and I would like to compute the matrix product of the i-th element in the first list with the i-th element in the second list. Of course, I could just loop over the lists:
for i in range(n): out[i] = dot( matrices1[i], matrices2[i] )
However, the list is quite long, and each matrix is very small (3x3) so this turns out to be quite slow. Is there a way to do this with a single numpy call? I have looked at tensordot but it outputs an N x N x 3 x 3 array, whereas I want an N x 3 x 3 output. I've also looked at various broadcasting tricks but I haven't found anything that works yet.
Alex _______________________________________________ NumPy-Discussion mailing list NumPy-Discussion@scipy.org http://mail.scipy.org/mailman/listinfo/numpy-discussion
Thank you.
On Sun, Aug 26, 2012 at 11:04 AM, Alex Flint alex.flint@gmail.com wrote:
I have two lists of 3x3 arrays and I would like to compute the matrix product of the i-th element in the first list with the i-th element in the second list. Of course, I could just loop over the lists:
for i in range(n): out[i] = dot( matrices1[i], matrices2[i] )
However, the list is quite long, and each matrix is very small (3x3) so this turns out to be quite slow. Is there a way to do this with a single numpy call? I have looked at tensordot but it outputs an N x N x 3 x 3 array, whereas I want an N x 3 x 3 output. I've also looked at various broadcasting tricks but I haven't found anything that works yet.
Alex
participants (2)
-
Alex Flint -
Jonathan Taylor