# Compute a series of matrix multiplications and matrix norms quickly in Python

I need to compute a series of matrix multiplications involving 3x3 matrices and a series of matrix norms also involving 3x3 matrices and I wonder how I can set these computations up with numpy such that they are computed quickly (preferably below 10ms).

I am given matrices $$A,B \in \mathbb{C}^{3 \times 3}$$ and I need to compute the following:

$$\tilde{R}_i=AC_iB$$

$$o_i=||\tilde{R}\tilde{R}^{T}-I_3||_F$$

Where $$\mathbb{C}_i \in \mathbb{C}^{3 \times 3}, i\in[1,...,N]$$ are matrices that are known beforehand and do not change.

My current approach looks like this:

import numpy as np
from numpy.linalg import norm
import time

#some constants
N=2**14
C=np.random.rand(3,3,N)+1j*np.random.rand(3,3,N)

#function that does computations
def compute_o(A,B):
global C
#perform the computations
####################################
T=np.einsum("ij,jkl",A,C)
R=np.einsum("lkn,kj",T,B)
####################################

#generate random inputs
A=np.random.rand(3,3)+1j*np.random.rand(3,3)
B=np.random.rand(3,3)+1j*np.random.rand(3,3)

#perform computations and estimate time
tic=time.time()
compute_o(A,B)
toc=time.time()
print("Time ellapsed: %fms"%(1e3*(toc-tic)))


Basically, I did the initial batch of matrix multiplications using einsum which on my computer takes around 2ms. However, I don't quite understand how I can compute the second step without using some sort of loops. Does anyone have an idea?

You can use einsum with explicit output indices, running over dimension n without summing.
T = np.einsum("ij,jkn", A, C)

• You could also use S=np.matmul(R,R,axes=[(0,1),(1,0),(0,1)]) - np.eye(3).reshape((3,3,1)) instead of the einsum, but surprisingly this turned out to perform slightly worse for me. No idea why, though I suspect somewhere there's a different loop order which isn't as cache friendly since the results do differ at floating point precision. Apr 23, 2023 at 15:23