fix test

Author: 6Ulm

ot/unbalanced.py

@@ -443,6 +443,8 @@ 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, :]
@@ -652,6 +654,8 @@ 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

@@ -104,16 +104,16 @@ def test_unbalanced_warmstart(nx, method, reg_type):
         reg_type=reg_type, warmstart=None, verbose=True
     )
 
-    # dim_a, dim_b = M.shape
-    # warmstart = (nx.zeros(dim_a, type_as=M), nx.zeros(dim_b, type_as=M))
-    # G, log = ot.unbalanced.sinkhorn_unbalanced(
-    #     a, b, M, reg=epsilon, reg_m=reg_m, method=method,
-    #     reg_type=reg_type, warmstart=warmstart, log=True, verbose=True
-    # )
-    # loss = ot.unbalanced.sinkhorn_unbalanced2(
-    #     a, b, M, reg=epsilon, reg_m=reg_m, method=method,
-    #     reg_type=reg_type, warmstart=warmstart, verbose=True
-    # )
+    dim_a, dim_b = M.shape
+    warmstart = (nx.zeros(dim_a, type_as=M), nx.zeros(dim_b, type_as=M))
+    G, log = ot.unbalanced.sinkhorn_unbalanced(
+        a, b, M, reg=epsilon, reg_m=reg_m, method=method,
+        reg_type=reg_type, warmstart=warmstart, log=True, verbose=True
+    )
+    loss = ot.unbalanced.sinkhorn_unbalanced2(
+        a, b, M, reg=epsilon, reg_m=reg_m, method=method,
+        reg_type=reg_type, warmstart=warmstart, verbose=True
+    )
 
     _, log = ot.lp.emd(a, b, M, log=True)
     warmstart1 = (log["u"], log["v"])
@@ -126,19 +126,19 @@ 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(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(
-    #     nx.to_numpy(log["logv"]), nx.to_numpy(log0["logv"]), atol=1e-05)
+    np.testing.assert_allclose(
+        nx.to_numpy(log["logu"]), nx.to_numpy(log0["logu"]), atol=1e-05)
+    np.testing.assert_allclose(
+        nx.to_numpy(log["logv"]), nx.to_numpy(log0["logv"]), atol=1e-05)
     np.testing.assert_allclose(
         nx.to_numpy(log0["logu"]), nx.to_numpy(log1["logu"]), atol=1e-05)
     np.testing.assert_allclose(
         nx.to_numpy(log0["logv"]), nx.to_numpy(log1["logv"]), atol=1e-05)
 
-    # np.testing.assert_allclose(nx.to_numpy(G), nx.to_numpy(G0), atol=1e-05)
+    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)