replace reshape by numpy slicing in return

1 files changed, 10 insertions, 2 deletions

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