opt_einsum icon indicating copy to clipboard operation
opt_einsum copied to clipboard

upgrade optimal and greedy algorithms with versions from cotengra

Open jcmgray opened this issue 9 months ago • 2 comments

I wrote upgraded versions of both the 'dp' (i.e. optimal), greedy and random-greedy pathfinders for cotengra: https://github.com/jcmgray/cotengra/blob/main/cotengra/pathfinders/path_basic.py.

They have some main advantages:

  1. they perform various simplifications first, namely:
    • ignore any indices that appear in all terms
    • combine any repeated indices within a single term
    • reduce any non-output indices that only appear on a single term
    • combine any scalar terms
    • combine any tensors with matching indices (hadamard products)
  2. they are faster thanks to a tweak of how the indices are counted .
  3. they have drop-in rust versions in https://github.com/jcmgray/cotengrust that are even faster
  4. the random greedy algorithm produces better paths by sampling an 'alpha' parameter.

Some disadvantages:

  • no memory_limit support
  • probably some other minor options dropped for random-greedy.

Including dropping the current optimal (non-dp) implementation these should fix #99, #112, #114, #167, #189, #243, #248, #233, #94.

Raising this issue to:

  1. gauge interest in porting these over to opt_einsum to replace greedy, dp, optimal and possibly random-greedy,
  2. see whether e.g. changes like dropping the memory_limit kwarg would be acceptable.
  3. see if anyone would like to take on this task!

jcmgray avatar Mar 11 '25 01:03 jcmgray

I wonder if we can get these to be an exact drop in replacement. A few thoughts:

  • Can we add memory_limit support?
  • For 1) how do we view the resulting changes in contraction behavior?

I like the rust idea and have been thinking about how to add it to opt_einsum proper. I think I can work through the shipping issues and have a reference and rust implementation side by side defaulting to the reference implementation if the rust library isn't available.

dgasmith avatar Jun 14 '25 15:06 dgasmith

Can we add memory_limit support?

I think one possible option here, that would make life easier, would be to simply write a separate function that that runs after the main path finding, and takes whatever path was returned and truncates it according to memory_limit. I think that would largely yield the same effect? and wouldn't require baking support in to every path finder.

For 1) how do we view the resulting changes in contraction behavior?

For the simplifications, these are returned as single item contractions e.g. path=[(1,), (0,), (0, 1)]. Might need to check if they will be handled properly by the contractor but I think they should already.

I like the rust idea and have been thinking about how to add it to opt_einsum proper. I think I can work through the shipping issues and have a reference and rust implementation side by side defaulting to the reference implementation if the rust library isn't available.

Yes how I have cotengra and cotengrust set up is that the functions in cotengrust exactly mirror the api of existing python implementations in cotengra, so it is a purely optional acceleration. Shipping via pypi is very easy with pyo3 and maturin, I didn't figure out conda-forge yet however.

jcmgray avatar Jun 26 '25 18:06 jcmgray