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

Empty step in contractions with scalar by dp #187

Open Geositta2000 opened 2 years ago

Geositta2000 commented 2 years ago

For the following code

from opt_einsum import contract
import numpy as np
import opt_einsum as oe
import timeit

dim = 30
I = np.random.rand(dim, dim, dim, dim)
C = np.random.rand(dim, dim)

path_info = oe.contract_expression(', pi,qj,ijkl,rk,sl->pqrs', (), (dim, dim), (dim, dim), (dim, dim, dim, dim), (dim, dim), (dim, dim))
print(path_info)

path_info = oe.contract_expression(', pi,qj,ijkl,rk,sl->pqrs', (), (dim, dim), (dim, dim), (dim, dim, dim, dim), (dim, dim), (dim, dim), optimize = 'dp')
print(path_info)

I got

<ContractExpression(', pi,qj,ijkl,rk,sl->pqrs')>
  1.  'ijkl,pi->jklp' [GEMM]
  2.  'jklp,qj->klpq' [GEMM]
  3.  'klpq,rk->lpqr' [GEMM]
  4.  'lpqr,sl->pqrs' [GEMM]
  5.  'pqrs,->pqrs' [OUTER/EINSUM]
<ContractExpression(', pi,qj,ijkl,rk,sl->pqrs')>
  1.  '->'
  2.  'sl,ijkl->sijk' [GEMM]
  3.  'sijk,rk->sijr' [GEMM]
  4.  'sijr,qj->sirq' [TDOT]
  5.  'sirq,pi->srqp' [TDOT]
  6.  'srqp,->pqrs' [OUTER/EINSUM]

Seems dp optimization will create an empty step, 1. '->'. The timing does not matter too much

 %timeit oe.contract(',pi,qj,ijkl,rk,sl->pqrs', 3.0,  C, C, I, C, C)
36.9 ms ± 1.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

 %timeit oe.contract(',pi,qj,ijkl,rk,sl->pqrs', 3.0,  C, C, I, C, C, optimize = 'dp')
38.2 ms ± 1.69 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Not sure if it holds in general.

So is 1. '->' from dp a feature?

dgasmith commented 2 years ago

The dp algorithm must identify that tensor as a disconnected subgraph to solve first. This is an oversight of the algorithm, but is unlikely to create a meaningful time discrepancy. We can leave this open to solve at some point.