I was just looking into numpy.einsum and encountered an issue which might
be worth pointing out in the documentation.
Let us say you wish to evaluate something like this (repeated indices a
D[alpha, alphaprime] = A[alpha, beta, sigma] * B[alphaprime, betaprime,
sigma] * C[beta, betaprime]
with einsum as
einsum('abs,cds,bd->ac', A, B, C)
then it is not exactly clear which order einsum evaluates the contractions
(or if it does it in one go). This can be very important since you can do
it in several ways, one of which has the least computational complexity.
The most efficient way of doing it is to contract e.g. A and C and then
contract that with B (or exchange A and B). A quick test on my labtop says
2.6 s with einsum and 0.13 s for two tensordots with A and B being D x D x
2 and C is D x D for D = 96. This scaling seems to explode for higher
dimensions, whereas it is much better with the two independent contractions
(i believe it should be O(D^3)).For D = 512 I could do it in 5 s with two
contractions, whereas I stopped waiting after 60 s for einsum (i guess
einsum probably is O(D^4) in this case).
I had in fact thought of making a function similar to einsum for a while,
but after I noticed it dropped it. I think, however, that there might still
be room for a tool for evaluating more complicated expressions efficiently.
I think the best way would be for the user to enter an expression like the
one above which is then evaluated in the optimal order. I know how to do
this (theoretically) if all the repeated indices only occur twice (like the
expression above), but for the more general expression supported by einsum
I om not sure how to do it (haven't thought about it). Here I am thinking
about stuff like x[i] = a[i] * b[i] and their more general counterparts (at
first glance this seems to be a simpler problem than full contractions). Do
you think there is a need/interest for this kind of thing? In that case I
would like the write it / help write it. Much of it, I think, can be
reduced to decomposing the expression into existing numpy operations (e.g.
tensordot). How to incorporate issues of storage layout etc, however, I
have no idea.
In any case I think it might be nice to write explicitly how the expression
in einsum is evaluated in the docs.
Department of Physics and Astronomy