[Numpy-discussion] einsum evaluation order
mwwiebe at gmail.com
Tue Jan 24 18:49:05 EST 2012
On Tue, Jan 24, 2012 at 6:32 AM, Søren Gammelmark <gammelmark at gmail.com>wrote:
> Dear all,
> 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).
You are correct, einsum presently uses the most naive evaluation.
> 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.
I think a good approach would be to modify einsum so it decomposes the
expression into multiple products. It may even just be a simple dynamic
programming problem, but I haven't given it much thought.
In any case I think it might be nice to write explicitly how the expression
> in einsum is evaluated in the docs.
That's a good idea, yes.
> Søren Gammelmark
> Department of Physics and Astronomy
> Aarhus University
> NumPy-Discussion mailing list
> NumPy-Discussion at scipy.org
-------------- next part --------------
An HTML attachment was scrubbed...
More information about the NumPy-Discussion