Vectorize per-label computations for sinkhorn l1l2

Author: kachayev

ot/backend.py

@@ -407,7 +407,7 @@ def power(self, a, exponents):
         """
         raise NotImplementedError()
 
-    def norm(self, a):
+    def norm(self, a, axis=None):
         r"""
         Computes the matrix frobenius norm.
 
@@ -627,7 +627,7 @@ def diag(self, a, k=0):
         """
         raise NotImplementedError()
 
-    def unique(self, a):
+    def unique(self, a, return_inverse=False):
         r"""
         Finds unique elements of given tensor.
 
@@ -1087,8 +1087,8 @@ def sqrt(self, a):
     def power(self, a, exponents):
         return np.power(a, exponents)
 
-    def norm(self, a):
-        return np.sqrt(np.sum(np.square(a)))
+    def norm(self, a, axis=None):
+        return np.linalg.norm(a, axis=axis)
 
     def any(self, a):
         return np.any(a)
@@ -1164,8 +1164,8 @@ def meshgrid(self, a, b):
     def diag(self, a, k=0):
         return np.diag(a, k)
 
-    def unique(self, a):
-        return np.unique(a)
+    def unique(self, a, return_inverse=False):
+        return np.unique(a, return_inverse=return_inverse)
 
     def logsumexp(self, a, axis=None):
         return special.logsumexp(a, axis=axis)
@@ -1461,8 +1461,8 @@ def sqrt(self, a):
     def power(self, a, exponents):
         return jnp.power(a, exponents)
 
-    def norm(self, a):
-        return jnp.sqrt(jnp.sum(jnp.square(a)))
+    def norm(self, a, axis=None):
+        return jnp.linalg.norm(a, axis=axis)
 
     def any(self, a):
         return jnp.any(a)
@@ -1535,8 +1535,8 @@ def meshgrid(self, a, b):
     def diag(self, a, k=0):
         return jnp.diag(a, k)
 
-    def unique(self, a):
-        return jnp.unique(a)
+    def unique(self, a, return_inverse=False):
+        return jnp.unique(a, return_inverse=return_inverse)
 
     def logsumexp(self, a, axis=None):
         return jspecial.logsumexp(a, axis=axis)
@@ -1881,8 +1881,8 @@ def sqrt(self, a):
     def power(self, a, exponents):
         return torch.pow(a, exponents)
 
-    def norm(self, a):
-        return torch.sqrt(torch.sum(torch.square(a)))
+    def norm(self, a, axis=None):
+        return torch.linalg.norm(a, dim=axis)
 
     def any(self, a):
         return torch.any(a)
@@ -1986,8 +1986,8 @@ def meshgrid(self, a, b):
     def diag(self, a, k=0):
         return torch.diag(a, diagonal=k)
 
-    def unique(self, a):
-        return torch.unique(a)
+    def unique(self, a, return_inverse=False):
+        return torch.unique(a, return_inverse=return_inverse)
 
     def logsumexp(self, a, axis=None):
         if axis is not None:
@@ -2306,8 +2306,8 @@ def power(self, a, exponents):
     def dot(self, a, b):
         return cp.dot(a, b)
 
-    def norm(self, a):
-        return cp.sqrt(cp.sum(cp.square(a)))
+    def norm(self, a, axis=None):
+        return cp.linalg.norm(a, axis=axis)
 
     def any(self, a):
         return cp.any(a)
@@ -2383,8 +2383,8 @@ def meshgrid(self, a, b):
     def diag(self, a, k=0):
         return cp.diag(a, k)
 
-    def unique(self, a):
-        return cp.unique(a)
+    def unique(self, a, return_inverse=False):
+        return cp.unique(a, return_inverse=return_inverse)
 
     def logsumexp(self, a, axis=None):
         # Taken from
@@ -2717,8 +2717,8 @@ def sqrt(self, a):
     def power(self, a, exponents):
         return tnp.power(a, exponents)
 
-    def norm(self, a):
-        return tf.math.reduce_euclidean_norm(a)
+    def norm(self, a, axis=None):
+        return tf.math.reduce_euclidean_norm(a, axis=axis)
 
     def any(self, a):
         return tnp.any(a)
@@ -2790,8 +2790,10 @@ def meshgrid(self, a, b):
     def diag(self, a, k=0):
         return tnp.diag(a, k)
 
-    def unique(self, a):
-        return tf.sort(tf.unique(tf.reshape(a, [-1]))[0])
+    def unique(self, a, return_inverse=False):
+        y, idx = tf.unique(tf.reshape(a, [-1]))
+        sort_idx = tf.argsort(y)
+        return y[sort_idx] if not return_inverse else (y[sort_idx], idx[sort_idx])
 
     def logsumexp(self, a, axis=None):
         return tf.math.reduce_logsumexp(a, axis=axis)

ot/da.py

@@ -235,16 +235,14 @@ def sinkhorn_l1l2_gl(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
     a, labels_a, b, M = list_to_array(a, labels_a, b, M)
     nx = get_backend(a, labels_a, b, M)
 
-    lstlab = nx.unique(labels_a)
-
+    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):
-        res = 0
-        for i in range(G.shape[1]):
-            for lab in lstlab:
-                temp = G[labels_a == lab, i]
-                res += nx.norm(temp)
-        return res
+        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

@@ -800,3 +800,30 @@ def test_emd_laplace_class(nx):
     transp_ys = otda.inverse_transform_labels(yt)
     assert_equal(transp_ys.shape[0], ys.shape[0])
     assert_equal(transp_ys.shape[1], len(np.unique(nx.to_numpy(yt))))
+
+
+def test_sinkhorn_l1l2_gl_cost_vectorized():
+    n_samples, n_labels = 150, 3
+    rng = np.random.RandomState(42)
+    G = rng.rand(n_samples, n_samples)
+    labels_a = rng.randint(n_labels, size=(n_samples,))
+
+    # previously used implementation for the cost estimator
+    lstlab = np.unique(labels_a)
+    def f(G):
+        res = 0
+        for i in range(G.shape[1]):
+            for lab in lstlab:
+                temp = G[labels_a == lab, i]
+                res += np.linalg.norm(temp)
+        return res
+
+    # new vectorized implementation for the cost estimator
+    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))