numpy: efficient sum of kronecker products

Question

• I have three sets of matrices {A_i}, {B_i}, and {C_i} with n matrices in each set
• The A_i are of dimension l x m, the B_i are of dimension m x o and the C_i are of dimension p x q
• I would like to compute the following:

Here is a concrete example for what I am after

A = np.arange(12).reshape(2,3,2)
B = np.arange(12,24).reshape(2,2,3)
C = np.arange(32).reshape(2,4,4)

result = np.zeros((12,12))
for i in range(2):
    result += np.kron(A[i,:,:] @ B[i,:,:], C[i,:,:])

How can I implement this more efficiently?

Many thanks for your help!

Answer 1

As suggested, I had a look into numpy.einsum. This turned out to be quite nice. A solution is:

np.einsum('ijk,imn->jmkn', np.einsum('ijk,ikm->ijm', A, B), C).reshape(A.shape[1] * C.shape[1], B.shape[2] * C.shape[2])

• The inner np.einsum() produces a 3d array of the products of the 2d "slices" of A and B
• The outer np.einsum() mimics (after appropriate reshaping) the kronecker product of this 3d matrix and C and summation.

I found the following two posts very helpful:

• Understanding NumPy's einsum
• Speeding Up Kronecker Products Numpy