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

contract_path with optimize='optimal' not working on small example #243

Open ryan112358 opened 1 week ago

ryan112358 commented 1 week ago

Cross posting from https://github.com/jax-ml/jax/issues/24929

import opt_einsum
import numpy as np

formula = 'a,c,d,db,ab,cb,ac,cd,ad,b->dbc'

arrays = [np.random.rand(*(2,)*len(key)) for key in formula.split('->')[0].split(',')]
opt_einsum.contract_path(
    formula, *arrays, einsum_call=True, use_blas=True, optimize='optimal')  # this hangs and does not complete
dgasmith commented 6 days ago

One caution is the optimal algorithm scales factorially, with 10 input arrays you are generally pushing the limits with 10's of seconds to minutes runtime. You will find similar results with the dp algorithm which has much better runtimes or this particular contraction can be evaluated with greedy for what is likely to be the optimal path. We generally recommend optimize='auto' to balance path finding/path quality for input formula. More information can be found here.

I'll check out your example to see if it's an actual hang or just a long path evaluation time.