update sliced plans with sparse matrix for tf compatibility

Author: lchapel

ot/sliced.py

@@ -825,7 +825,8 @@ def sliced_plans(
             X_theta, Y_theta, a, b, p, require_sort=True, return_plan=True
         )
 
-        plan = plan.tocsr().tocoo()  # especially for tensorflow compatibility
+        if str(nx) == "tensorflow":  # tf does not support duplicate entries
+            plan = [plan[k].tocsr().tocoo() for k in range(n_proj)]
 
         if str(nx) == "jax":
             plan = [nx.todense(plan[k]) for k in range(n_proj)]
@@ -853,26 +854,30 @@ def sliced_plans(
                 )
                 for k in range(n_proj)
             ]
-
-        if metric in ("minkowski", "euclidean", "cityblock"):
-            costs = [
-                nx.sum(
-                    (
-                        (nx.sum(nx.abs(X[plan[k].row] - Y[plan[k].col]) ** p, axis=1))
-                        ** (1 / p)
+        else:
+            if metric in ("minkowski", "euclidean", "cityblock"):
+                costs = [
+                    nx.sum(
+                        (
+                            (
+                                nx.sum(
+                                    nx.abs(X[plan[k].row] - Y[plan[k].col]) ** p, axis=1
+                                )
+                            )
+                            ** (1 / p)
+                        )
+                        * plan[k].data
                     )
-                    * plan[k].data
-                )
-                for k in range(n_proj)
-            ]
-        else:  # metric = "sqeuclidean"
-            costs = [
-                nx.sum(
-                    (nx.sum((X[plan[k].row] - Y[plan[k].col]) ** 2, axis=1))
-                    * plan[k].data
-                )
-                for k in range(n_proj)
-            ]
+                    for k in range(n_proj)
+                ]
+            else:  # metric = "sqeuclidean"
+                costs = [
+                    nx.sum(
+                        (nx.sum((X[plan[k].row] - Y[plan[k].col]) ** 2, axis=1))
+                        * plan[k].data
+                    )
+                    for k in range(n_proj)
+                ]
 
     if dense:
         plan = [nx.todense(plan[k]) for k in range(n_proj)]

test/test_sliced.py

@@ -776,7 +776,7 @@ def test_sliced_plans():
     ot.sliced.sliced_plans(x, y, thetas=thetas, metric="minkowski")
 
     # test with an unsupported metric
-    with pytest.raises(ValueError):
+    with pytest.raises(AssertionError):
         ot.sliced.sliced_plans(x, y, thetas=thetas, metric="mahalanobis")
 
     # test with a warm theta
@@ -798,7 +798,7 @@ def test_sliced_plans():
     b /= b.sum()
 
     # test with the minkowski metric
-    ot.sliced.sliced_plans(x, y, metric="minkowski")
+    ot.sliced.sliced_plans(x, y, n_proj=10, metric="minkowski")
 
 
 def test_min_pivot_sliced():
@@ -857,7 +857,7 @@ def test_min_pivot_sliced():
     ot.sliced.min_pivot_sliced(x, y, thetas=thetas, metric="cityblock")
 
     # test with an unsupported metric
-    with pytest.raises(ValueError):
+    with pytest.raises(AssertionError):
         ot.sliced.min_pivot_sliced(x, y, thetas=thetas, metric="mahalanobis")
 
     # test with a warm theta
@@ -922,7 +922,7 @@ def test_expected_sliced():
     ot.sliced.expected_sliced(x, y, thetas=thetas, metric="minkowski")
 
     # test with an unsupported metric
-    with pytest.raises(ValueError):
+    with pytest.raises(AssertionError):
         ot.sliced.expected_sliced(x, y, thetas=thetas, metric="mahalanobis")