dgasmith / opt_einsum

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

Way to keep tensor Ids from being rearranged to cretae custom paths #206

Closed LionFrangoulis closed 1 year ago

LionFrangoulis commented 1 year ago

According to the example in The Optimal Path example , upon contracting a pair of tensors, the resulting one will be pushed to the end of the list. Is there a way to prevent this (putting the resulting tensor on the position of the first contracted one), so I can create custom paths through my network?

jcmgray commented 1 year ago

Hi @LionFrangoulis, yes the actual raw path format is a slightly hard to reason with directly. Your best bet is just to construct it however you like (e.g. the 'ssa' format), then convert to the direct path format at the end.

I might also suggest cotengra if you want to really explicitly construct the contraction tree manually and visualize it etc - https://cotengra.readthedocs.io/en/latest/trees.html. It can produce the opt_einsum compatible path among other formats.

LionFrangoulis commented 1 year ago

@jcmgray Thank you very much, once I understood the ssa format, evreything worked out just fine.