Speedup Sinkhorn with einsum + bench

[rflamary]

Hello,

Since @LeoGautheron and the sklearn devellopers opened my eyes about einsum, I found some new speedup especially on the test of violation of the marginals constraints. I also use einsum to compute the final loss for paralell sinkhorn computations.

With the following benchmark

import time
import numpy as np
import ot
rng = np.random.RandomState(0)

n=5000
d=2
b=5
reg=1e0
Xs, ys, Xt = rng.randn(n, d), rng.randint(0, 2, size=n), rng.randn(n, d)

a=ot.unif(n)
B=np.ones((n,b))/n

M=ot.dist(Xs,Xt)

time1 = time.time()
res=ot.sinkhorn(a,a,M,reg,verbose=True)
time2 = time.time()
print("Sinkhorn Computation Time {:6.2f} sec".format(time2-time1))

time1 = time.time()
res=ot.sinkhorn(a,B,M,reg,verbose=True)
time2 = time.time()
print("Multi Sinkhorn Computation Time {:6.2f} sec".format(time2-time1))
print(res)

I get the following outputs for master and my branch

Branch master
python3 mytest/test_einsum.py
It.  |Err         
-------------------
    0|5.449272e-06|
   10|7.115862e-12|
Sinkhorn Computation Time   2.07 sec
It.  |Err         
-------------------
    0|1.293152e+00|
   10|6.770311e-07|
   20|2.770657e-11|
Multi Sinkhorn Computation Time   8.75 sec
[0.8783663 0.8783663 0.8783663 0.8783663 0.8783663]
Branch speedup
python3 mytest/test_einsum.py
It.  |Err         
-------------------
    0|5.449272e-06|
   10|7.115862e-12|
Sinkhorn Computation Time   1.73 sec
It.  |Err         
-------------------
    0|1.293152e+00|
   10|6.770311e-07|
   20|2.770657e-11|
Multi Sinkhorn Computation Time   8.08 sec
[0.8783663 0.8783663 0.8783663 0.8783663 0.8783663]

I also added a nice make option for benchmarks that run a script twice by comparing it to the master branch with the following command:

make bench script=yourscript.py branch=thebranch

[LeoGautheron]

nice !