From 92b7075568207a468cc821cd6a21e130b9d89f96 Mon Sep 17 00:00:00 2001
From: "Mokhtar Z. Alaya" <mokhtarzahdi.alaya@gmail.com>
Date: Tue, 7 Jan 2020 16:21:40 +0100
Subject: replace reshape by numpy slicing in return

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

(limited to 'ot/bregman.py')

diff --git a/ot/bregman.py b/ot/bregman.py
index 8a20307..8cfea7e 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -2107,6 +2107,14 @@ def screenkhorn(a, b, M, reg, ns_budget=None, nt_budget=None, uniform=True, rest
     vsc_full = np.full(nt, epsilon * kappa)
     usc_full[I] = usc
     vsc_full[J] = vsc
+    
+    if log:
+      log['u'] = usc_full
+      log['v'] = vsc_full
+      
+    gamma = usc_full[:, None] * K * vsc_full[None, :]
 
-    Gsc = usc_full.reshape((-1, 1)) * K * vsc_full.reshape((1, -1))
-    return Gsc
+    if log:
+      return gamma, log
+    else:
+      return gamma
-- 
cgit v1.2.3