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

Question

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?

Divakar · Accepted Answer · 2017-09-25 04:33:25Z

2

We could use np.dot, like so -

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

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

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

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

np.matmul(A,B)

answered Sep 25, 2017 at 4:28

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Or in python > 3.5, A @ B