⬆️ Update dependencies

yhs0602 · Jan 18, 2025 · d9de8e3 · d9de8e3
1 parent 5917d97
commit d9de8e3
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -36,11 +36,10 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-jax = [
-    "jax>=0.4.10",
-    "jaxlib>=0.4.10; platform_system != 'Darwin'",  # CUDA/CPU backend (Not Apple GPU)
-    "jax-metal>=0.1.1; platform_system == 'Darwin'",  # Metal backend (Apple GPU)
-]
+jax_cuda11 = ["jax[cuda11]"]
+jax_cuda12 = ["jax[cuda12]"]
+jax_cpu = ["jax[cpu]"]
+jax_metal = ["jax-metal"]
 
 [project.urls]
 Source = "https://github.com/yhs0602/CraftGround"