From f4bfeb73da098384aa67599e7f729fb683a1bcc9 Mon Sep 17 00:00:00 2001
From: Rémi Flamary <remi.flamary@gmail.com>
Date: Tue, 24 Jul 2018 15:54:56 +0200
Subject: ensum tets marginals sinkhorn

---
 ot/bregman.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

(limited to 'ot')

diff --git a/ot/bregman.py b/ot/bregman.py
index 58e74de..c755f51 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -396,10 +396,7 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000,
         log['v'] = v
 
     if nbb:  # return only loss
-        res = np.zeros((nbb))
-        for i in range(nbb):
-            res[i] = np.sum(
-                u[:, i].reshape((-1, 1)) * K * v[:, i].reshape((1, -1)) * M)
+        res = np.einsum('ik,ij,jk,ij->k', u, K, v, M)
         if log:
             return res, log
         else:
-- 
cgit v1.2.3