Merge branch 'installation_refactor'

dattalab · Oct 18, 2024 · eff3453 · eff3453
2 parents 4499238 + c9b3359
commit eff3453
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -14,11 +14,21 @@ python_requires = >=3.8
 install_requires =
     numba>=0.56.4
     dynamax>=0.1.2
+    tqdm
+    optree
+
+[options.extras_require]
+cuda11 = 
+    jax==0.3.22
+    etils==1.5.2
+    scipy==1.11.3
     chex==0.1.6
     jaxtyping==0.2.14
     tensorflow_probability==0.19.0
-    tqdm
-    optree
+
+cuda12 =
+    jax>=0.4.16
+    scipy==1.12.0
 
 [options.package_data]
 * = *.md