CODEWITHSUNDEEP
pythonnumpymatrix-multiplicationnumpy-ndarray
arundathi-c
arundathi-c

Reputation: 15

Multiply a 3D matrix with 2D matrix to give a 2D matrix

I have 2 matrices. Say A of size 6610 and B of size 610. What I want to do is multiply the 10 66 matrices of A with the 10 6 element vectors of B to give a 6*10 matrix. Is there a way to do this without using a loop?

What I want is

A = np.ones((6,6,10))
B = np.ones((6,10))
mat = np.zeros((6,10))
for i in range(10):
    mat[:,i] = A[:,:,i]@B[:,i]

but without the for loop.

Upvotes: 1

Views: 55

Answers (1)

Divakar
Divakar

Reputation: 221574

We can use np.einsum -

mat = np.einsum('ijk,jk->ik',A,B)

Alternatively, with np.matmul/@-operator -

mat = (A.transpose(2,0,1)@B.T[:,:,None])[...,0].T

Upvotes: 2

Related Questions