Black voodoo of NumPy Einsum

Question

I got some working code using einsum function. But as einsum is currently still like black voodoo for me. I was wondering, what this code actually is doing and if it can be somehow optimized using np.dot

My data looks likes this

n, p, q = 40000, 8, 4
a = np.random.rand(n, p, q)
b = np.random.rand(n, p)

And my existing functions einsum functions looks like this

f1 = np.einsum("ijx,ijy->ixy", a, a)
f2 = np.einsum("ijx,ij->ix", a, b)

But what does it really do? I get till here: each dimension (axis) is represented by a label, i is equal to the first axis n, j for the 2nd axis p and x and y are different labels for the same axis q. So the order of the output array of f1 is ixy and thus the output shape is 40000,4,4 (n,q,q)

But that's as far as I get. And

Answer 1

Lets play around with a couple of small arrays

In [110]: a=np.arange(2*3*4).reshape(2,3,4)

In [111]: b=np.arange(2*3).reshape(2,3)

In [112]: np.einsum('ijx,ij->ix',a,b)
Out[112]: 
array([[ 20,  23,  26,  29],
       [200, 212, 224, 236]])

In [113]: np.diagonal(np.dot(b,a)).T
Out[113]: 
array([[ 20,  23,  26,  29],
       [200, 212, 224, 236]])

np.dot operates on the last dim of the 1st array, and 2nd to the last of the 2nd. So I have to switch the arguments so the 3 dimension lines up. dot(b,a) produces a (2,2,4) array. diagonal selects 2 of those 'rows', and transpose to clean up. Another einsum expresses that cleanup nicely:

In [122]: np.einsum('iik->ik',np.dot(b,a))

Since np.dot is producing a larger array than the original einsum, it is unlikely to be faster, even if the underlying C code is tighter.

(Curiously I'm having trouble replicating np.dot(b,a) with einsum; it won't generate that (2,2,...) array).

For the a,a case we have to do something similar - roll the axes of one array so the last dimension lines up with the 2nd to last of the other, do the dot, and then cleanup with diagonal and transpose:

In [157]: np.einsum('ijx,ijy->ixy',a,a).shape
Out[157]: (2, 4, 4)
In [158]: np.einsum('ijjx->jix',np.dot(np.rollaxis(a,2),a))
In [176]: np.diagonal(np.dot(np.rollaxis(a,2),a),0,2).T

tensordot is another way of taking a dot over selected axes.

np.tensordot(a,a,(1,1))
np.diagonal(np.rollaxis(np.tensordot(a,a,(1,1)),1),0,2).T  # with cleanup