diff options
Diffstat (limited to 'ot/bregman.py')
| -rw-r--r-- | ot/bregman.py | 11 |
1 files changed, 8 insertions, 3 deletions
diff --git a/ot/bregman.py b/ot/bregman.py index d2ade46..29ca9fd 100644 --- a/ot/bregman.py +++ b/ot/bregman.py @@ -358,9 +358,14 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000, while (err > stopThr and cpt < numItermax): uprev = u vprev = v - KtransposeU = np.einsum('ij,i->j',K,u)#np.dot(K.T, u) - v = np.divide(b, KtransposeU) - u = 1. / np.einsum('ij,j->i',Kp,v)#np.dot(Kp, v) + if nbb: + KtransposeU = np.einsum('ij,i,k->jk',K,u)#np.dot(K.T, u) + v = np.divide(b, KtransposeU) + u = 1. / np.einsum('ij,jk->ik',Kp,v)#np.dot(Kp, v) + else: + KtransposeU = np.einsum('ij,i->j',K,u)#np.dot(K.T, u) + v = np.divide(b, KtransposeU) + u = 1. / np.einsum('ij,j->i',Kp,v)#np.dot(Kp, v) if (np.any(KtransposeU == 0) or np.any(np.isnan(u)) or np.any(np.isnan(v)) or |