depend on jax backend in extra

lenskit · May 28, 2024 · c3f60d7 · c3f60d7
1 parent c25a5e8
commit c3f60d7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -57,7 +57,7 @@ torch = ["torch"]
 cupy = ["cupy"]
 cuda11x = ["cupy-cuda11x"]
 tf = ["tensorflow >=2,<3"]
-jax = ["jax>=0.4.16"]
+jax = ["jax[cpu]>=0.4.16"]
 
 [project.urls]
 Homepage = "https://seedbank.lenksit.org"