summaryrefslogtreecommitdiff
path: root/ot/bregman.py
diff options
context:
space:
mode:
authorRémi Flamary <remi.flamary@gmail.com>2017-01-05 13:40:38 +0100
committerRémi Flamary <remi.flamary@gmail.com>2017-01-05 13:40:38 +0100
commitbb6020dd69933735a6cf821eabcbedca92f2ecb4 (patch)
treeead967f53cec3a0a70f8adc14707a47a86cc6c3e /ot/bregman.py
parent17315cb08ebf777626dc9aec4bf140cfb5c42503 (diff)
more efficient sinkhorn
Diffstat (limited to 'ot/bregman.py')
-rw-r--r--ot/bregman.py4
1 files changed, 2 insertions, 2 deletions
diff --git a/ot/bregman.py b/ot/bregman.py
index 568ead2..fd22259 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -141,9 +141,9 @@ def sinkhorn(a,b, M, reg, numItermax = 1000, stopThr=1e-9, verbose=False, log=Fa
cpt = cpt +1
#print 'err=',err,' cpt=',cpt
if log:
- return np.dot(np.diag(u),np.dot(K,np.diag(v))),log
+ return u.reshape((-1,1))*K*v.reshape((1,-1)),log
else:
- return np.dot(np.diag(u),np.dot(K,np.diag(v)))
+ return u.reshape((-1,1))*K*v.reshape((1,-1))
def sinkhorn_stabilized(a,b, M, reg, numItermax = 1000,tau=1e3, stopThr=1e-9,warmstart=None, verbose=False,print_period=20, log=False):
"""