CODEWITHSUNDEEP
pythonnumpymatrix
secondrate
secondrate

Reputation: 149

Multiply i-th 2-d matrix in numpy 3d array with i-th column in 2d array

Suppose that I have a 3d array A and a 2d array B. A has dimension (s,m,m) while B has dimension (m,s).

I want to write code for a 2d array C with dimension (m,s) such that C[:,i] = A[i,:,:] @ B[:,i].

Is there a way to do this elegantly without using a for loop in numpy?

One solution I thought of was to reshape B into a 3d array with dimension (m,s,1), multiply A and B via A@B, then reshape the resulting 3d array into a 2d array. This sounds a bit tedious and was wondering if tensordot or einsum can be applied here.

Suggestions appreciated. Thanks!

Upvotes: 0

Views: 74

Answers (1)

Paul Panzer
Paul Panzer

Reputation: 53029

Using einsum is straight forward here:

A = np.arange(18).reshape(2,3,3)
B = np.arange(6).reshape(3,2)

C = np.einsum("ijk,ki->ji",A,B)

for i in range(2):
   A[i]@B[:,i]==C[:,i]

# array([ True,  True,  True])
# array([ True,  True,  True])

Upvotes: 1

Related Questions