fix jax versions

mlcommons · Feb 8, 2025 · 39bb876 · 39bb876
1 parent 082be03
commit 39bb876
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -103,13 +103,13 @@ jax_core_deps = [
   "protobuf==4.25.5",
 ]
 jax_cpu = [
-  "jax==0.4.26",
-  "jaxlib==0.4.26",
+  "jax==0.4.28",
+  "jaxlib==0.4.28",
   "algorithmic_efficiency[jax_core_deps]",
 ]
 jax_gpu = [
-  "jax==0.4.26",
-  "jaxlib==0.4.26",
+  "jax==0.4.28",
+  "jaxlib==0.4.28",
   "jax-cuda12-plugin[with_cuda]==0.4.28",
   "jax-cuda12-pjrt==0.4.28",
   "algorithmic_efficiency[jax_core_deps]",