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 ellipsis fails when `shapes=True` #235

Closed nova77 closed 4 days ago

nova77 commented 3 weeks ago

Easy way to reproduce:

input = np.random.rand(1,3,5,2,)
a = np.random.rand(2)
b = np.random.rand(7)

# ValueError: operands could not be broadcast together with shapes (4,) (0,) 
oe.contract_path('...D,D,k->...k', input.shape, a.shape, b.shape, shapes=True)

While if we do not pass the shapes, the function works:

oe.contract_path('...D,D,k->...k', input, a, b)  # all good!

The source of error is in in parser.py where operands (i.e. the shapes) are converted into arrays:

operands = [possibly_convert_to_numpy(x) for x in operands[1:]]

The fix should be simple: just skip the conversion when shapes is True.

nova77 commented 3 weeks ago

The smallest change would be:

    if shapes:
        operand_shapes = [list(s) for s in operands]  # from operand_shapes = operands
    else:
        operand_shapes = [o.shape for o in operands]