dgasmith / opt_einsum

⚡️Optimizing einsum functions in NumPy, Tensorflow, Dask, and more with contraction order optimization.
https://dgasmith.github.io/opt_einsum/
MIT License
863 stars 68 forks source link

Error when using PyTorch backend on large contraction #92

Closed emprice closed 4 years ago

emprice commented 5 years ago

Great project! I'm looking forward to using it in my research.

I came across this when I was trying to optimize a tensor (as in the quimb documentation) and wanted to try my hand at a contraction manually. My machine was using 88GB RAM to perform the contraction, which is simply too much, and I wanted to see which step of the contraction was causing a problem, since none of my tensors are very large.

When I tried to put together a MWE, everything seems to work great until I actually want to do the contraction, using the PyTorch backend. It's important for my application to use PyTorch because I need to do the optimization step (are there other backends that would let me do that?) for a time-evolved state. I get the error: RuntimeError: only lowercase letters a-z allowed as indices.

Now, this is obviously a limitation in PyTorch itself, not opt_einsum, but I'm wondering if there might be a clever workaround that I've missed? I could imagine that it would be possible to take a pre-computed contraction path and

  1. Limit the number of indices in any one step (maybe not, now that I think about it more), and/or
  2. Cleverly re-index so that separate calls to einsum use indices that start over at "a", which might mitigate or eliminate the problem, depending on how many indices are contracted in a given step

If it helps, this is the contraction I was trying, which I generated for a 2-dimensional quimb TensorNetwork:

abc,dbef,gehi,jhkl,mkn,aopq,drpst,gusvw,jxvyz,mAyB,oCDE,rFDGH,uIGJK,xLJMN,AOMP,CQRS,FTRUV,IWUXY,LZXÀÁ,OÂÀÃ,QÄÅ,TÄÆÇ,WÆÈÉ,ZÈÊË,ÂÊÌ,ÍÎcÏ,ÐÎÑfÒ,ÓÑÔiÕ,ÖÔ×lØ,Ù×nÚ,ÍÛÜqÝ,ÐÞÜßtà,Óáßâwã,Öäâåzæ,ÙçåBè,ÛéêEë,ÞìêíHî,áïíðKñ,äòðóNô,çõóPö,é÷øSù,ìúøûVü,ïýûþYÿ,òĀþāÁĂ,õăāÃĄ,÷ąÅĆ,úąćÇĈ,ýćĉÉĊ,ĀĉċËČ,ăċÌč,ĎďÏ,ĐďđÒ,ĒđēÕ,ĔēĕØ,ĖĕÚ,ĎėĘÝ,ĐęĘĚà,ĒěĚĜã,ĔĝĜĞæ,ĖğĞè,ėĠġë,ęĢġģî,ěĤģĥñ,ĝĦĥħô,ğĨħö,ĠĩĪù,ĢīĪĬü,ĤĭĬĮÿ,ĦįĮİĂ,ĨıİĄ,ĩIJĆ,īIJijĈ,ĭijĴĊ,įĴĵČ,ıĵč->

Thanks in advance for any help! Depending on the feasibility of the above, I might be able to help out in developing functionality that would allow contractions like this one, since they're important for my project.

jcmgray commented 5 years ago

Hi @emprice. This is sadly a limitation on the pytorch end. To clarify with regards to your points:

  1. Trying to get the number of indices to a minimum on intermediate tensors (which can be much larger than the inputs) during a contraction is the problem of trying to find a efficient contraction path essentially! There is something called tensor network slicing where you sum over certain indices, reducing the memory a bit, but the more indices you slice over the more time the contraction takes (if you sliced over all indices it would be like a naive einsum so ~10^44 times slower in this case!).
  2. opt_einsum already internally converts indices into the a-z range preferably, but if you check contract_path the problem is that in your case an intermediate tensor has more than 26 indices - there just aren't enough lowercase letters

Some ways forward:

dgasmith commented 5 years ago

@emprice Great that you are adding this to pytorch directly!

dgasmith commented 4 years ago

@emprice See here for some ideas on slicing which would help. Good luck on the PyTorch integration!

Going to close this for now, let us know if we can help in the future.

jcmgray commented 4 years ago

You can also try slicing here - will hopefully get time to add to opt_einsum at some point soon.