Skip to content

Commit 5e3392a

Browse files
committed
cancel einsum
1 parent 603c0ee commit 5e3392a

File tree

1 file changed

+5
-8
lines changed

1 file changed

+5
-8
lines changed

ot/bregman.py

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -358,14 +358,11 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000,
358358
while (err > stopThr and cpt < numItermax):
359359
uprev = u
360360
vprev = v
361-
if nbb:
362-
KtransposeU = np.einsum('ij,ik->jk',K,u)#np.dot(K.T, u)
363-
v = np.divide(b, KtransposeU)
364-
u = 1. / np.einsum('ij,jk->ik',Kp,v)#np.dot(Kp, v)
365-
else:
366-
KtransposeU = np.einsum('ij,i->j',K,u)#np.dot(K.T, u)
367-
v = np.divide(b, KtransposeU)
368-
u = 1. / np.einsum('ij,j->i',Kp,v)#np.dot(Kp, v)
361+
362+
KtransposeU = np.dot(K.T, u)
363+
v = np.divide(b, KtransposeU)
364+
u = 1. / np.dot(Kp, v)
365+
369366

370367
if (np.any(KtransposeU == 0) or
371368
np.any(np.isnan(u)) or np.any(np.isnan(v)) or

0 commit comments

Comments
 (0)