CODEWITHSUNDEEP
pythonnumpy
martianwars
martianwars

Reputation: 6500

Matrix multiplication of inner dimensions of 3D tensors?

I have two numpy matrices of dimension (386, 3, 4) and (386, 4, 3). I want to produce an output dimension of (386, 3, 3). In other words, I wish to execute the following loop in a vectorized fashion -

for i in range(len(input1)):
    output[i] = np.matmul(input1[i], input2[i])

What's the best way to do this?

Upvotes: 2

Views: 1385

Answers (2)

Divakar
Divakar

Reputation: 221534

We need to keep the first axes aligned, so I would suggest using an approach with np.einsum -

np.einsum('ijk,ikl->ijl',input1,input2)

Sample run to verify shapes -

In [106]: a = np.random.rand(386, 3, 4)

In [107]: b = np.random.rand(386, 4, 3)

In [108]: np.einsum('ijk,ikl->ijl',a,b).shape
Out[108]: (386, 3, 3)

Sample run to verify values on smaller input -

In [174]: a = np.random.rand(2, 3, 4)

In [175]: b = np.random.rand(2, 4, 3)

In [176]: output = np.zeros((2,3,3))

In [177]: for i in range(len(a)):
     ...:     output[i] = np.matmul(a[i], b[i])
     ...:     

In [178]: output
Out[178]: 
array([[[ 1.43473795,  0.860279  ,  1.17855877],
        [ 1.91036828,  1.23063125,  1.5319063 ],
        [ 1.06489098,  0.86868941,  0.84986621]],

       [[ 1.07178572,  1.020091  ,  0.63070531],
        [ 1.34033495,  1.26641131,  0.79911685],
        [ 1.68916831,  1.63009854,  1.14612462]]])

In [179]: np.einsum('ijk,ikl->ijl',a,b)
Out[179]: 
array([[[ 1.43473795,  0.860279  ,  1.17855877],
        [ 1.91036828,  1.23063125,  1.5319063 ],
        [ 1.06489098,  0.86868941,  0.84986621]],

       [[ 1.07178572,  1.020091  ,  0.63070531],
        [ 1.34033495,  1.26641131,  0.79911685],
        [ 1.68916831,  1.63009854,  1.14612462]]])

Sample run to verify values on bigger input -

In [180]: a = np.random.rand(386, 3, 4)

In [181]: b = np.random.rand(386, 4, 3)

In [182]: output = np.zeros((386,3,3))

In [183]: for i in range(len(a)):
     ...:     output[i] = np.matmul(a[i], b[i])
     ...:     

In [184]: np.allclose(np.einsum('ijk,ikl->ijl',a,b), output)
Out[184]: True

Upvotes: 2

Paul Panzer
Paul Panzer

Reputation: 53029

matmul also works:

a = np.random.random((243,3,4))
b = np.random.random((243,4,3))
np.matmul(a,b).shape
# (243, 3, 3)

Upvotes: 4

Related Questions