python - Fast elementwise multiplication of a tensor by a list of vectors

Question

I have a tensor and I want to multiply that tensor into a list of vectors. An example of minimal code is below.

tensor=np.arange(4*5*6).reshape(4,5,6)
vectorList=[]
vectorList.append(np.array([0,1,2,3,4,5]))
vectorList.append(np.array([6,7,8,9,10,11]))
vectorList.append(np.array([12,13,14,15,16,17]))

results=[]
for i in vectorList:
  results.append(tensor*i)

This produces the required results but that for loop is just to slow. I can't seem to find a numpy/tensorflow elementwise multiply command that will allow this multiplcation after a stacking of the vectorList.

I'm willing to trade memory for speed if that is an issue. I'm also unconcern about the container for the results. A list of tensors or a 4 dimensional tensor are fine here.

Edit: Updated the vectors in the list to make the question more clear.

question from:https://stackoverflow.com/questions/66052253/fast-elementwise-multiplication-of-a-tensor-by-a-list-of-vectors

Answer 1

Your 3d array (renamed from tensor):

In [448]: arr.shape
Out[448]: (4, 5, 6)

vectosList as array is 2d:

In [449]: np.array(vectorList).shape
Out[449]: (3, 6)

And your results, as array, is 4d:

In [450]: np.array(results).shape
Out[450]: (3, 4, 5, 6)

Thanks for giving a nice diverse set of dimensions. It's easier to track them that way.

We can produce the same 4d array with broadcasting. I've included all the ':' just to highlight how dimensions are paired:

In [451]: res = np.array(vectorList)[:,None,None,:]*arr[None,:,:,:]
In [452]: res.shape
Out[452]: (3, 4, 5, 6)
In [453]: np.allclose(res, np.array(results))
Out[453]: True

res=np.array(vectorList)[:,None,None]*arr is the same thing.

A deleted answer suggested einsum, the correct expression is np.einsum('il,jkl->ijkl',vectorList, arr).