fix: dependency list

EMI-Group · Apr 16, 2024 · 4bd3f66 · 4bd3f66
1 parent 488fac5
commit 4bd3f66
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -43,6 +43,7 @@ test = [
   "chex >= 0.1.0",
   "flax >= 0.5.0",
   "pytest >= 6.0.0",
+  "tensorflow >= 2.12.0",
 ]
 
 vis = [
@@ -69,6 +70,11 @@ full = [
   "gpjax >= 0.8.0",
   "plotly >= 5.0.0",
   "pandas >= 2.0.0",
+  "tensorflow-datasets >= 4.0.0",
+  "grain >= 0.1.0",
+  "brax >= 0.1.0",
+  "plotly >= 5.0.0",
+  "pandas >= 2.0.0",
 ]
 
 gp = ["gpjax >= 0.8.0"]