update sliced plans with sparse matrix for tf compatibility

Author: lchapel

ot/sliced.py

@@ -873,21 +873,24 @@ def sliced_plans(
                             )
                             ** (1 / p)
                         )
-                        * plan[k].data
+                        * plan[k]
                     )
                     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
+                        * plan[k]
                     )
                     for k in range(n_proj)
                 ]
 
     if dense and not str(nx) == "jax":
         plan = [nx.todense(plan[k]) for k in range(n_proj)]
+    elif str(nx) == "jax":
+        warnings.warn("JAX does not support sparse matrices, converting to dense")
+        plan = [nx.todense(plan[k]) for k in range(n_proj)]
 
     if log:
         log_dict = {"X_theta": X_theta, "Y_theta": Y_theta, "thetas": thetas}