CODEWITHSUNDEEP
pythonnumpymatrix-multiplication
dshin
dshin

Reputation: 2398

Dot product for arbitrary shaped NumPy arrays

Given two numpy.ndarray objects, A and B, of arbitrary shapes, I'd like to compute an numpy.ndarray C with the property that C[i] == np.dot(A[i], B[i]) for all i. How can I do this?

Example 1: A.shape==(2,3,4) and B.shape==(2,4,5), then we should have C.shape==(2,3,5).

Example 2: A.shape==(2,3,4) and B.shape==(2,4), then we should have C.shape==(2,3).

Upvotes: 3

Views: 220

Answers (2)

Divakar
Divakar

Reputation: 221514

Here's a generic solution to cover all kinds of cases / arbitrary shapes using some reshaping and np.einsum. einsum helps here as we need alignment along the first axis and reduction along the last axes of the input arrays. The implementation would look something like this -

def dotprod_axis0(A,B):
    N,nA,nB = A.shape[0], A.shape[-1], B.shape[1]
    Ar = A.reshape(N,-1,nA)
    Br = B.reshape(N,nB,-1)
    return np.squeeze(np.einsum('ijk,ikl->ijl',Ar,Br))

Cases

I. A : 2D, B : 2D

In [119]: # Inputs
     ...: A = np.random.randint(0,9,(3,4))
     ...: B = np.random.randint(0,9,(3,4))
     ...: 

In [120]: for i in range(A.shape[0]):
     ...:     print np.dot(A[i], B[i])
     ...:     
33
86
48

In [121]: dotprod_axis0(A,B)
Out[121]: array([33, 86, 48])

II. A : 3D, B : 3D

In [122]: # Inputs
     ...: A = np.random.randint(0,9,(2,3,4))
     ...: B = np.random.randint(0,9,(2,4,5))
     ...: 

In [123]: for i in range(A.shape[0]):
     ...:     print np.dot(A[i], B[i])
     ...:     
[[ 74  70  53 118  43]
 [ 47  43  29  95  30]
 [ 41  37  26  23  15]]
[[ 50  86  33  35  82]
 [ 78 126  40 124 140]
 [ 67  88  35  47  83]]

In [124]: dotprod_axis0(A,B)
Out[124]: 
array([[[ 74,  70,  53, 118,  43],
        [ 47,  43,  29,  95,  30],
        [ 41,  37,  26,  23,  15]],

       [[ 50,  86,  33,  35,  82],
        [ 78, 126,  40, 124, 140],
        [ 67,  88,  35,  47,  83]]])

III. A : 3D, B : 2D

In [125]: # Inputs
     ...: A = np.random.randint(0,9,(2,3,4))
     ...: B = np.random.randint(0,9,(2,4))
     ...: 

In [126]: for i in range(A.shape[0]):
     ...:     print np.dot(A[i], B[i])
     ...:     
[ 87 105  53]
[152 135 120]

In [127]: dotprod_axis0(A,B)
Out[127]: 
array([[ 87, 105,  53],
       [152, 135, 120]])

IV. A : 2D, B : 3D

In [128]: # Inputs
     ...: A = np.random.randint(0,9,(2,4))
     ...: B = np.random.randint(0,9,(2,4,5))
     ...: 

In [129]: for i in range(A.shape[0]):
     ...:     print np.dot(A[i], B[i])
     ...:     
[76 93 31 75 16]
[ 33  98  49 117 111]

In [130]: dotprod_axis0(A,B)
Out[130]: 
array([[ 76,  93,  31,  75,  16],
       [ 33,  98,  49, 117, 111]])

Upvotes: 4

user2357112
user2357112

Reputation: 280227

Assuming you want ordinary matrix multiplication for dot (not, say, matrix-vector or the weird crap dot does for higher dimensions), then sufficiently recent NumPy versions (1.10+) let you do

C = numpy.matmul(A, B)

and sufficiently recent Python versions (3.5+) let you write that as

C = A @ B

assuming your NumPy is also sufficiently recent.

Upvotes: 0

Related Questions