According to this, to do a matrix multiplication with NumPy’s einsum(), the call would be indeed np.einsum("ij, jk -> ik", A, B) supposing A is nx3 and B is 3x3. Notice the two ‘k’ trailing, not following the usual matrix multiplication notation. I guess that was a coding mistake that went unnoticed/not tested.