CODEWITHSUNDEEP
pythonarraysnumpy
beedrill
beedrill

Reputation: 413

NumPy dot prodoct of a 2d array and a 3D array element wise

A = np.array([[1,1],
              [2,2],
              [3,3]])

B = np.array([[[1],[2]],[[3],[4]]])

I think of B as an array of 2 matrix, what I want to achieve is to do dot product between A and each element of B, to get:

[[[3],
  [6],
  [9]], 
 [[7],
  [14],
  [21]]]

but if I do np.dot(A,B), I get 

[[3,7],
 [6,14],
 [9,21]]

how to get what I want here?

Upvotes: 2

Views: 1873

Answers (1)

Divakar
Divakar

Reputation: 221764

We could use np.dot, like so -

A.dot(B).T[0,...,None]

Or with np.tensordot -

np.tensordot(B[...,0], A, axes=((1),(1)))[...,None]

Or np.einsum -

np.einsum('ijk,lj->ilk',B,A)

np.matmul seems to be working as well without any additional work -

np.matmul(A,B)

Upvotes: 2

Related Questions