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

contract_path with `optimize=bool` raises #219

Open ev-br opened 10 months ago

ev-br commented 10 months ago

contract_path's optimize argument is documented as str, list or bool, but using bool arguments raises a TypeError:

n [1]: import opt_einsum as oe

In [2]: import numpy as np

In [3]: a = np.ones((2, 2))

In [4]: oe.contract_path('ij, jk, kl', a, a, a, optimize=False)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[4], line 1
----> 1 oe.contract_path('ij, jk, kl', a, a, a, optimize=False)

File ~/.conda/envs/pytorch-dev/lib/python3.8/site-packages/opt_einsum/contract.py:276, in contract_path(*operands, **kwargs)
    273 contraction_list = []
    275 # Build contraction tuple (positions, gemm, einsum_str, remaining)
--> 276 for cnum, contract_inds in enumerate(path):
    277     # Make sure we remove inds from right to left
    278     contract_inds = tuple(sorted(list(contract_inds), reverse=True))
    280     contract_tuple = helpers.find_contraction(contract_inds, input_sets, output_set)

TypeError: 'bool' object is not iterable

IIUC this is due to https://github.com/dgasmith/opt_einsum/blob/master/opt_einsum/contract.py#L305C1-L307C25:

    if not isinstance(path_type, (str, paths.PathOptimizer)):
        # Custom path supplied
        path = path_type

Granted, this is an edge case and maybe just a documentation issue of not allowing bools in the first place.

jcmgray commented 10 months ago

I suppose the contract_path for optimize=False is a bit ill-defined. The options would be:

  1. Simply update the docs
  2. Always return path = [range(num_inputs)], which would allow contract_path to be called with the same args as contract always - probably preferable.
ev-br commented 10 months ago

What about contract_path=True, should this be equivalent to "auto"?

dgasmith commented 9 months ago

I would agree contract_path=True should be equivalent to auto. It would seem natively to me that contract_path=False is equivalent to ensue or path = [range(num_inputs)].