带有矩阵数组的Numpy矩阵乘法

Question

I have several numpy arrays that I would like to multiply (using dot , so matrix multiplication). I'd like to put them all into a numpy array, but I can't figure out how to do it.

Eg

a = np.random.randn((10,2,2))
b = np.random.randn((10,2))

So I have 10 2x2 matrices (a) and 10 2x1 matrices (b). What I could do is this:

c = np.zeros((10,2))
for i in range(10):
  c[i] = np.dot(a[i,:,:],b[i,:])

You get the idea.

But I feel like there's a usage of dot or tensordot or something that would do this in one line really easily. I just can't make sense of the dot and tensordot functions for >2 dimensions like this.

Answer 1

You could use np.einsum :

c = np.einsum('ijk,ik->ij', a, b)

einsum performs a sum of products. Since matrix multiplication is a sum of products, any matrix multiplication can be expressed using einsum . It is based on Einstein summation notation .

The first argument to einsum, ijk,ik->ij is a string of subscripts . ijk declares that a has three axes which are denoted by i , j , and k .

ik , similarly, declares that the axes of b will be denoted i and k .

When the subscripts repeat, those axes are locked together for purposes of summation. The part of the subscript that follows the -> shows the axes which will remain after summation.

Since the k appears on the left (of the -> ) but disappears on the right, there is summation over k . It means that the sum

c_ij = sum over k ( a_ijk * b_ik )

should be computed. Since this sum can be computed for each i and j , the result is an array with subscripts i and j .