CODEWITHSUNDEEP
pythonnumpymatrixtheano.scan
Phong Khac Do
Phong Khac Do

Reputation: 53

3D Matrix multiplication in numpy and theano

I have a matrix A with size (5,7,3) and a matrix B with size (5,3,8). I want to multiply them C = A.B, and the size of C is (5,7,8).

It means that one 2D submatrix with size (7,3) in matrix A will be multiplied with one 2D submatrix with size (3,8) in matrix B respectively. So I have to multiply 5 times.

The simplest way is using a loop and numpy:

for u in range(5):
    C[u] = numpy.dot(A[u],B[u])

Is there any way to do this without using a loop? Is there any equivalent method in Theano to do this without using scan?

Upvotes: 2

Views: 475

Answers (1)

Daniel F
Daniel F

Reputation: 14399

Can be done pretty simply with np.einsum in numpy.

C = numpy.einsum('ijk,ikl->ijl', A, B)

It can also simply be:

C = numpy.matmul(A,B)

Since the docs state:

If either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly

Theano has similar functionaly of batched_dot so it would be

C = theano.tensor.batched_dot(A, B)

Upvotes: 3

Related Questions