Fix flake

Author: kachayev

ot/da.py

@@ -238,11 +238,13 @@ def sinkhorn_l1l2_gl(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
     labels_u, labels_idx = nx.unique(labels_a, return_inverse=True)
     n_labels = labels_u.shape[0]
     unroll_labels_idx = nx.eye(n_labels, type_as=labels_u)[None, labels_idx]
+
     def f(G):
         G_split = nx.repeat(G.T[:, :, None], n_labels, axis=2)
         return nx.norm(G_split * unroll_labels_idx, axis=1).sum()
 
     lstlab = nx.unique(labels_a)
+
     def df(G):
         W = nx.zeros(G.shape, type_as=G)
         for i in range(G.shape[1]):

test/test_da.py

@@ -810,6 +810,7 @@ def test_sinkhorn_l1l2_gl_cost_vectorized():
 
     # previously used implementation for the cost estimator
     lstlab = np.unique(labels_a)
+
     def f(G):
         res = 0
         for i in range(G.shape[1]):
@@ -822,8 +823,9 @@ def f(G):
     lstlab, lstlab_idx = np.unique(labels_a, return_inverse=True)
     n_samples = lstlab.shape[0]
     midx = np.eye(n_samples, dtype='int32')[None, lstlab_idx]
+
     def f2(G):
         G_split = np.repeat(G.T[:, :, None], n_samples, axis=2)
         return np.linalg.norm(G_split * midx, axis=1).sum()
 
-    assert np.allclose(f(G), f2(G))
+    assert np.allclose(f(G), f2(G))