From 16f51f971607efab2c73958d207c582b389406c8 Mon Sep 17 00:00:00 2001
From: Leo gautheron <gautheron@iv-cm-359.creatis.insa-lyon.fr>
Date: Thu, 20 Apr 2017 12:12:15 +0200
Subject: sinkhorn GPU implementation

---
 ot/bregman.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

(limited to 'ot/bregman.py')

diff --git a/ot/bregman.py b/ot/bregman.py
index 0453f14..c46e5dc 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -112,9 +112,11 @@ def sinkhorn(a,b, M, reg, numItermax = 1000, stopThr=1e-9, verbose=False, log=Fa
     while (err>stopThr and cpt<numItermax):
         uprev = u
         vprev = v
-        v = np.divide(b,np.dot(K.T,u))
+        KtransposeU = np.dot(K.T, u)
+        v = np.divide(b, KtransposeU)
         u = 1./np.dot(Kp,v)
-        if (np.any(np.dot(K.T,u)==0) or
+
+        if (np.any(KtransposeU==0) or
            np.any(np.isnan(u)) or np.any(np.isnan(v)) or
            np.any(np.isinf(u)) or np.any(np.isinf(v))):
             # we have reached the machine precision
-- 
cgit v1.2.3