attempt to fix dtype bug

theislab · Nov 27, 2024 · 0c8be0f · 0c8be0f
1 parent 4080a94
commit 0c8be0f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -54,7 +54,7 @@ dependencies = [
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
     "docrep>=0.3.2",
-    "ott-jax[neural]>=0.4.6,<=0.4.8",
+    "ott-jax[neural]>=0.4.6",
     "cloudpickle>=2.2.0",
     "rich>=13.5",
     "docstring_inheritance>=2.0.0",

diff --git a/src/moscot/backends/ott/_utils.py b/src/moscot/backends/ott/_utils.py
@@ -132,7 +132,7 @@ def ensure_2d(arr: ArrayLike, *, reshape: bool = False) -> jax.Array:
         return jnp.reshape(arr, (-1, 1))
     if arr.ndim != 2:
         raise ValueError(f"Expected array to have 2 dimensions, found `{arr.ndim}`.")
-    return arr
+    return arr.astype(jnp.float64)
 
 
 def convert_scipy_sparse(arr: Union[sp.spmatrix, jesp.BCOO]) -> jesp.BCOO: