3d Matrix multiplication in numpy

2
votes

I am using numpy to perform matrix multiplication, and I cannot figure out how to leverage numpy for 3d matrix multiplication.

Say I have a 3x3 matrix, a, and I multiply it by a 3x1 vector, b. This will give a 3x1 vector, c.

This is done in numpy with:

# (3, 3) * (3, 1) -> (3, 1)
c = np.matmul(a, b)

Ok, so now I want to perform a similar operation on a 3d matrix that is essentially 2500 3x3 matrices. Right now I am doing something to the effect of:

# (2500, 3, 3) * (2500, 3, 1) -> list of (3, 1) vectors with length 2500
C = [np.matmul(a, b) for a, b in zip(A, B)]

which returns a list of (3, 1) vectors.

I would rather NOT loop and instead fully leverage numpy's vectorization and matrix/tensor products. Is there some operation so I can do...

# (2500, 3, 3) * (2500, 3, 1) -> (2500, 3, 1)
np.<function>(A, B, <args>)

I've seen stuff about using np.tensordot, but I don't know how to set axes.

np.tensordot(A, B, axes=???)
2
pythonnumpymatrixlinear-algebra

2 Answers

4
votes

For array of dimension 3 (or, a rank-3 tensor) that you have, you can use np.einsum doc for more complex matrix multiplications. In your particular case, you can use the following

>>> import numpy as np
>>> x = np.random.randint(0, 3, (3, 3, 3))  # shape (3, 3, 3)
>>> y = np.random.randint(0, 3, (3, 3, 3))  # shape (3, 3, 3)
>>> np.einsum('ijk,ikl->ijl', x, y)  # still shape (3, 3, 3)

In particular, the einsum expression 'ijk,ikl->ijl' means for each ith matrix, do a regular matrix multiplication jk,kl->jl and put the result in the ith entry in the resulting tensor (ndarray). A more general form of this process could be

np.einsum('...jk,...kl->...jl', x, y)

where you can have arbitrary number of dimensions in front of each tensor (ndarray) that you have.

See the following for a complete example:

>>> import numpy as np
>>> x = np.random.randint(0, 3, (3, 3, 3))  # shape (3, 3, 3)
>>> x
array([[[0, 0, 1],
        [2, 2, 1],
        [2, 1, 1]],

       [[2, 0, 2],
        [2, 2, 1],
        [2, 2, 2]],

       [[2, 2, 2],
        [1, 1, 2],
        [0, 2, 2]]])
>>> y = np.random.randint(0, 3, (3, 3, 3))  # shape (3, 3, 3)
>>> y
array([[[0, 0, 1],
        [2, 1, 0],
        [0, 0, 2]],

       [[1, 2, 0],
        [2, 0, 1],
        [2, 2, 1]],

       [[0, 2, 1],
        [0, 1, 0],
        [0, 2, 1]]])
>>> np.einsum('ijk,ikl->ijl', x, y)
array([[[ 0,  0,  2],
        [ 4,  2,  4],
        [ 2,  1,  4]],

       [[ 6,  8,  2],
        [ 8,  6,  3],
        [10,  8,  4]],

       [[ 0, 10,  4],
        [ 0,  7,  3],
        [ 0,  6,  2]]])
>>> np.einsum('...ij,...jk->...ik', x, y)
array([[[ 0,  0,  2],
        [ 4,  2,  4],
        [ 2,  1,  4]],

       [[ 6,  8,  2],
        [ 8,  6,  3],
        [10,  8,  4]],

       [[ 0, 10,  4],
        [ 0,  7,  3],
        [ 0,  6,  2]]])
1
votes

np.matmul(A,B) works just fine. What error(s) did you get?

In [263]: A,B = np.arange(24).reshape(2,3,4), np.arange(8).reshape(2,4,1)

The einsum solution:

In [264]: np.einsum('ijk,ikl->ijl',A,B)
Out[264]: 
array([[[ 14],
        [ 38],
        [ 62]],

       [[302],
        [390],
        [478]]])
In [265]: _.shape
Out[265]: (2, 3, 1)

The matmul solution:

In [266]: A@B
Out[266]: 
array([[[ 14],
        [ 38],
        [ 62]],

       [[302],
        [390],
        [478]]])

Your loop:

In [267]: [np.matmul(a, b) for a, b in zip(A, B)]
Out[267]: 
[array([[14],
        [38],
        [62]]),
 array([[302],
        [390],
        [478]])]

matmuldocs:

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.