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

reusing contraction generated from contract_path leads to 'ValueError: Number of einsum subscripts must be equal to the number of operands.' #164

Closed ghost closed 3 years ago

ghost commented 3 years ago

Hi, I tried to reuse path generated from contract_path. based on https://optimized-einsum.readthedocs.io/en/stable/reusing_paths.html https://optimized-einsum.readthedocs.io/en/stable/autosummary/opt_einsum.contract_path.html https://pypi.org/project/opt-einsum/

import numpy as np
from opt_einsum import contract, contract_path, contract_expression
import time

N = 10
C = np.random.rand(N, N)
I = np.random.rand(N, N, N, N)

start = time.time()
output= np.einsum('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C, optimize=True)
end = time.time()
print(end-start)

start = time.time()
output= contract('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C, optimize='optimal')
end = time.time()
print(end-start)

path, info = contract_path('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C, optimize="optimal")      
print(path,info)

my_expr = contract_path('pi,qj,ijkl,rk,sl->pqrs', path)

The above code leads to

ValueError: Number of einsum subscripts must be equal to the number of operands.

How to let it run normally? Thanks.

dgasmith commented 3 years ago

I believe you are looking for contract_expression like so:

expression = contract_expression(
    "pi,qj,ijkl,rk,sl->pqrs",
    C.shape,
    C.shape,
    I.shape,
    C.shape,
    C.shape,
    optimize="optimal",
)
print(expression)

start = time.time()
my_expr = expression(C, C, I, C, C)
end = time.time()
print(end - start)
ghost commented 3 years ago

Thank you so much! It works!