numpy 矩阵的向量化乘法

Question

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, eg 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

Answer 1

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)

Answer 2

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.