numpy Vectorized multiplication of matrices

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

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)