I'm struggling to understand exactly how einsum works. I've looked at the documentation and a few examples, but it's not seeming to stick.
Here's an example we went over in class:
C = np.einsum("ij,jk->ki", A, B)
for two arrays: A and B.
I think this would take A^T * B, but I'm not sure (it's taking the transpose of one of them right?). Can anyone walk me through exactly what's happening here (and in general when using einsum)?
8 Answers
(Note: this answer is based on a short blog post about einsum I wrote a while ago.)
What does einsum do?
Imagine that we have two multi-dimensional arrays, A and B. Now let's suppose we want to...
- multiply
AwithBin a particular way to create new array of products; and then maybe - sum this new array along particular axes; and then maybe
- transpose the axes of the new array in a particular order.
There's a good chance that einsum will help us do this faster and more memory-efficiently than combinations of the NumPy functions like multiply, sum and transpose will allow.