While working on a project I came across a situation where I had to do an element-wise multiplication of rows of two matrices to produce a 3rd-order tensor. A few minutes of Googling did not give me an efficient solution. So here is a solution that is faster than the naive approach of using python list comprehension.
Let be a matrix and be a matrix. We want a tensor such that , where is the i-th row of A and is the j-th row of B.
Assume we have the following matrices:
The following uses python list comprehension:
Solution using einsum
The following uses Einstein notation function einsum of numpy:
Not only is the einsum approach much shorter, it is also about 10 times faster than the naive approach. Here are some quick benchmark results on my system running OS X.
In the first approach creating a new array and reshaping it are additional operations that the einsum method doesn’t have. So to get an idea about the time it takes to just multiply the rows using python list comprehension I ran the following benchmark:
From the above it is clear that the bulk of the time is spent in looping, which is slow, rather than creating a numpy array and reshaping it.
The einsum function in numpy is a powerful construct that can be used to represent complex matrix operations in a compact way and can result in significant performance improvement over loops. It is worth learning the Einstein sum notation, even for it’s own sake — it is pretty cool after all.