Numpy: How to multiply (N,N) and (N,N,M,M) numpy arrays?

Question

I want to multiply two numpy arrays . One numpy array is given by matrix of shape (10, 10) and the other is given by a matrix of matrices, ie shape (10, 10, 256, 256) .

I now simply want to multiply each matrix in the second matrix of matrices with the corresponding component in the first matrix. For instance, the matrix at position (0, 0) in the second matrix shall be multiplied by the value at position (0, 0) in the first matrix.

Intuitively, this is not really complicated, but numpy does not seem to support that. Or at least I am not smart enough to make it work. The ValueError that is thrown says:

ValueError: operands could not be broadcast together with shapes (10,10) (10,10,256,256)

Can anybody of you help me please? How can I achieve what I want in a numpyy way.

Answer 1

You can use the NumPy einsum function, eg, (using zeros arrays as dummies in this example):

import numpy as np
x = np.zeros((10, 10))
y = np.zeros((10, 10, 256, 256))
z = np.einsum("ij,ijkm->km", x, y)
print(z.shape)
(256, 256)

See here for a nice description of einsum 's usage.