set ottjax version first and the tests

theislab · Mar 20, 2024 · fd407df · fd407df
1 parent a5187c0
commit fd407df
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -55,7 +55,7 @@ dependencies = [
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
     "docrep>=0.3.2",
-    "ott-jax>=0.4.5",
+    "ott-jax @ git+https://github.com/ott-jax/ott.git@51a658aaa026528a3b9702e3d25c8f7af598cef0",
     "cloudpickle>=2.2.0",
     "rich>=13.5",
 ]

diff --git a/tests/backends/ott/test_backend.py b/tests/backends/ott/test_backend.py
@@ -12,8 +12,8 @@
 from ott.problems.linear.linear_problem import LinearProblem
 from ott.problems.quadratic import quadratic_problem
 from ott.problems.quadratic.quadratic_problem import QuadraticProblem
+from ott.solvers.linear import solve as sinkhorn
 from ott.solvers.linear.sinkhorn import Sinkhorn
-from ott.solvers.linear.sinkhorn import solve as sinkhorn
 from ott.solvers.linear.sinkhorn_lr import LRSinkhorn
 from ott.solvers.quadratic.gromov_wasserstein import GromovWasserstein
 from ott.solvers.quadratic.gromov_wasserstein_lr import LRGromovWasserstein

diff --git a/tests/problems/base/test_compound_problem.py b/tests/problems/base/test_compound_problem.py
@@ -7,7 +7,7 @@
 import numpy as np
 from ott.geometry.costs import Cosine, Euclidean, SqEuclidean
 from ott.geometry.pointcloud import PointCloud
-from ott.solvers.linear.sinkhorn import solve as sinkhorn
+from ott.solvers.linear import solve as sinkhorn
 from sklearn.metrics.pairwise import euclidean_distances
 
 from anndata import AnnData

diff --git a/tests/problems/base/test_general_problem.py b/tests/problems/base/test_general_problem.py
@@ -6,7 +6,7 @@
 import numpy as np
 import pandas as pd
 from ott.geometry.pointcloud import PointCloud
-from ott.solvers.linear.sinkhorn import solve as sinkhorn
+from ott.solvers.linear import solve as sinkhorn
 
 import scanpy as sc
 from anndata import AnnData