Matrix multiplication with numpy.einsum

Question

I have the following two arrays with shape:

    A = (d,w,l)
    B = (d,q)

And I want to combine these into a 3d array with shape:

    C = (q,w,l)

To be a bit more specific, in my case d (depth of the 3d array) is 2, and i'd first like to multiply all positions out of w * l in the upper layer of A (so d = 0) with the first value of B in the highest row (so d=0, q=0). For d=1 I do the same, and then sum the two so:

    C_{q=0,w,l} = A_{d=0,w,l}*B_{d=0,q=0} + A_{d=1,w,l}*B_{d=1,q=0}

I wanted to calculate C by making use of numpy.einsum. I thought of the following code:

    A = np.arange(100).reshape(2,10,5)

    B = np.arange(18).reshape(2,9)

    C = np.einsum('ijk,i -> mjk',A,B)

Where ijk refers to 2,10,5 and mjk refers to 9,10,5. However I get an error. Is there some way to perform this multiplication with numpy einsum?

Thanks

Answer 1

Your shapes A = (d,w,l), B = (d,q), C = (q,w,l) practically write the einsum expression

C=np.einsum('dwl,dq->qwl',A,B)

which I can test with

In [457]: np.allclose(A[0,:,:]*B[0,0]+A[1,:,:]*B[1,0],C[0,:,:])
Out[457]: True