fix bug in test

Author: 6Ulm

ot/unbalanced.py

@@ -277,30 +277,55 @@ def sinkhorn_unbalanced2(a, b, M, reg, reg_m, method='sinkhorn',
     ot.unbalanced.sinkhorn_reg_scaling: Unbalanced Sinkhorn with epsilon scaling :ref:`[9, 10] <references-sinkhorn-unbalanced2>`
 
     """
-    b = list_to_array(b)
+    M, a, b = list_to_array(M, a, b)
+    nx = get_backend(M, a, b)
+
     if len(b.shape) < 2:
-        b = b[:, None]
+        if method.lower() == 'sinkhorn':
+            res = sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
+                                            warmstart, numItermax=numItermax,
+                                            stopThr=stopThr, verbose=verbose,
+                                            log=log, **kwargs)
+
+        elif method.lower() == 'sinkhorn_stabilized':
+            res = sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type,
+                                                 warmstart, numItermax=numItermax,
+                                                 stopThr=stopThr, verbose=verbose,
+                                                 log=log, **kwargs)
+        elif method.lower() in ['sinkhorn_reg_scaling']:
+            warnings.warn('Method not implemented yet. Using classic Sinkhorn-Knopp')
+            res = sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
+                                            warmstart, numItermax=numItermax,
+                                            stopThr=stopThr, verbose=verbose,
+                                            log=log, **kwargs)
+        else:
+            raise ValueError('Unknown method %s.' % method)
 
-    if method.lower() == 'sinkhorn':
-        return sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
-                                         warmstart, numItermax=numItermax,
-                                         stopThr=stopThr, verbose=verbose,
-                                         log=log, **kwargs)
+        if log:
+            return nx.sum(M * res[0]), res[1]
+        else:
+            return nx.sum(M * res)
 
-    elif method.lower() == 'sinkhorn_stabilized':
-        return sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type,
-                                              warmstart, numItermax=numItermax,
-                                              stopThr=stopThr,
-                                              verbose=verbose,
-                                              log=log, **kwargs)
-    elif method.lower() in ['sinkhorn_reg_scaling']:
-        warnings.warn('Method not implemented yet. Using classic Sinkhorn-Knopp')
-        return sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
-                                         warmstart, numItermax=numItermax,
-                                         stopThr=stopThr, verbose=verbose,
-                                         log=log, **kwargs)
     else:
-        raise ValueError('Unknown method %s.' % method)
+        if method.lower() == 'sinkhorn':
+            return sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
+                                             warmstart, numItermax=numItermax,
+                                             stopThr=stopThr, verbose=verbose,
+                                             log=log, **kwargs)
+
+        elif method.lower() == 'sinkhorn_stabilized':
+            return sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type,
+                                                  warmstart, numItermax=numItermax,
+                                                  stopThr=stopThr, verbose=verbose,
+                                                  log=log, **kwargs)
+        elif method.lower() in ['sinkhorn_reg_scaling']:
+            warnings.warn('Method not implemented yet. Using classic Sinkhorn-Knopp')
+            return sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
+                                             warmstart, numItermax=numItermax,
+                                             stopThr=stopThr, verbose=verbose,
+                                             log=log, **kwargs)
+        else:
+            raise ValueError('Unknown method %s.' % method)
 
 
 def sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type="entropy",
@@ -443,8 +468,6 @@ def sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type="entropy",
             v = nx.ones(dim_b, type_as=M)
     else:
         u, v = nx.exp(warmstart[0]), nx.exp(warmstart[1])
-        if not n_hists:
-            u, v = u.reshape(-1), v.reshape(-1)
 
     if reg_type == "kl":
         K = nx.exp(-M / reg) * a.reshape(-1)[:, None] * b.reshape(-1)[None, :]
@@ -654,8 +677,6 @@ def sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type="entropy",
             v = nx.ones(dim_b, type_as=M)
     else:
         u, v = nx.exp(warmstart[0]), nx.exp(warmstart[1])
-        if not n_hists:
-            u, v = u.reshape(-1), v.reshape(-1)
 
     if reg_type == "kl":
         log_ab = nx.log(a + 1e-16).reshape(-1)[:, None] + nx.log(b + 1e-16).reshape(-1)[None, :]

test/test_unbalanced.py

@@ -15,7 +15,6 @@
 
 
 @pytest.mark.parametrize("method,reg_type", itertools.product(["sinkhorn", "sinkhorn_stabilized"], ["kl", "entropy"]))
-# @pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_stabilized"])
 def test_unbalanced_convergence(nx, method, reg_type):
     # test generalized sinkhorn for unbalanced OT
     n = 100
@@ -80,7 +79,6 @@ def test_unbalanced_convergence(nx, method, reg_type):
 
 
 @pytest.mark.parametrize("method,reg_type", itertools.product(["sinkhorn", "sinkhorn_stabilized"], ["kl", "entropy"]))
-# @pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_stabilized"])
 def test_unbalanced_warmstart(nx, method, reg_type):
     # test generalized sinkhorn for unbalanced OT
     n = 100
@@ -115,8 +113,8 @@ def test_unbalanced_warmstart(nx, method, reg_type):
         reg_type=reg_type, warmstart=warmstart, verbose=True
     )
 
-    _, log = ot.lp.emd(a, b, M, log=True)
-    warmstart1 = (log["u"], log["v"])
+    _, log_emd = ot.lp.emd(a, b, M, log=True)
+    warmstart1 = (log_emd["u"], log_emd["v"])
     G1, log1 = ot.unbalanced.sinkhorn_unbalanced(
         a, b, M, reg=epsilon, reg_m=reg_m, method=method,
         reg_type=reg_type, warmstart=warmstart1, log=True, verbose=True
@@ -126,9 +124,6 @@ def test_unbalanced_warmstart(nx, method, reg_type):
         reg_type=reg_type, warmstart=warmstart1, verbose=True
     )
 
-    np.testing.assert_allclose(nx.to_numpy(loss), nx.to_numpy(loss0), atol=1e-5)
-    np.testing.assert_allclose(nx.to_numpy(loss0), nx.to_numpy(loss1), atol=1e-5)
-
     np.testing.assert_allclose(
         nx.to_numpy(log["logu"]), nx.to_numpy(log0["logu"]), atol=1e-05)
     np.testing.assert_allclose(
@@ -141,6 +136,9 @@ def test_unbalanced_warmstart(nx, method, reg_type):
     np.testing.assert_allclose(nx.to_numpy(G), nx.to_numpy(G0), atol=1e-05)
     np.testing.assert_allclose(nx.to_numpy(G0), nx.to_numpy(G1), atol=1e-05)
 
+    np.testing.assert_allclose(nx.to_numpy(loss), nx.to_numpy(loss0), atol=1e-5)
+    np.testing.assert_allclose(nx.to_numpy(loss0), nx.to_numpy(loss1), atol=1e-5)
+
 
 @pytest.mark.parametrize("method,reg_m", itertools.product(["sinkhorn", "sinkhorn_stabilized"], [1, float("inf")]))
 def test_unbalanced_relaxation_parameters(nx, method, reg_m):