numpy Vectorized multiplication of matrices-CodePudding

So what Im trying to do is this...

Suppose we have a set of matrices:

a1 = [1, 2]
b1 = [[3, 4], [5, 6]]
c1 = [7, 8]

a2 = [2, 4]
b2 = [[6, 8], [10, 12]]
c2 = [14, 16]

a3 = [3, 5]
b3 = [[7, 9], [11, 13]]
c3 = [15, 17]

A1 = np.array(a1)
B1 = np.array(b1)
C1 = np.array(c1)

A2 = np.array(a2)
B2 = np.array(b2)
C2 = np.array(c2)

A3 = np.array(a3)
B3 = np.array(b3)
C3 = np.array(c3)

I vectorized the collection of those arrays to produce

A = np.array([A1, A2, A3])
B = np.array([B1, B2, B3])
C = np.array([C1, C2, C3])

Then I need to perform something like

A.T @ B @ C

And I'd like to get the restult in the form of

np.array([A1.T @ B1 @ C1, A2.T @ B2 @ C2, A3.T @ B3 @ C3])

The actual number of these triplets of matrices is in hundreds of thousands, so I'd like to do it using numpy vectorization, e.g. some form of tensor multiplication

I've tried to apply np.tensordot, but it doesn't qutie work as expected:

np.tensordot(A, B, axes=((1), (1)))

Produces:

array([[**[13, 16]**,
        [26, 32],
        [29, 35]],

       [[26, 32],
        **[52, 64]**,
        [58, 70]],

       [[34, 42],
        [68, 84],
        **[76, 92]**]])

Whereas I'd like it to be just an array of those marked arrays (each corresponding to Ai.T @ Bi)

The simple A @ B produces the same result...

So is there a way to produce the expected result without a loop over first dime

CodePudding user response：

The shapes:

In [272]: A.shape
Out[272]: (3, 2)
In [273]: B.shape
Out[273]: (3, 2, 2)
In [274]: C.shape
Out[274]: (3, 2)

the subarray product is a scalar:

In [275]: (A1.T@B1@C1).shape
Out[275]: ()
In [276]: (A1.T@B1@C1)
Out[276]: 219

matmul/@ can treat the first dimension (3), as a batch one if we expand all arrays to 3d, keeping in mind that we want the "the last of A with the 2nd to the last B" rule the last 2 dimensions:

In [277]: A[:,None,:]@B@C[:,:,None]
Out[277]: 
array([[[ 219]],

       [[1752]],

       [[2704]]])
In [278]: _.shape
Out[278]: (3, 1, 1)
In [279]: (A[:,None,:]@B@C[:,:,None])[:,0,0]
Out[279]: array([ 219, 1752, 2704])

the equivalent einsum is:

In [280]: np.einsum('ij,ijk,ik->i', A, B, C)
Out[280]: array([ 219, 1752, 2704])

Here batch i is repeated all the way, while j and k are the familiar matrix product variables.

CodePudding user response：

Hi you can vectorize your problem adding some dimensions and doing the product:

import numpy as np


a1 = [1, 2]
b1 = [[3, 4], [5, 6]]
c1 = [7, 8]

a2 = [2, 4]
b2 = [[6, 8], [10, 12]]
c2 = [14, 16]

a3 = [3, 5]
b3 = [[7, 9], [11, 13]]
c3 = [15, 17]

A1 = np.array(a1)
B1 = np.array(b1)
C1 = np.array(c1)

A2 = np.array(a2)
B2 = np.array(b2)
C2 = np.array(c2)

A3 = np.array(a3)
B3 = np.array(b3)
C3 = np.array(c3)

A = np.array([A1, A2, A3])
B = np.array([B1, B2, B3])
C = np.array([C1, C2, C3])

Res = (A[:,:,None]*B*C[:,None,:]).sum((1,2))

EDIT

You also can solve the same problem using Einsten sum notation (i've just realized that). This is more straight forward and also easier.

res = np.einsum('ni,nij,nj->n',A,B,C)