pin jax to 0.6.2

priyakasimbeg · priyakasimbeg · commit db652ee84b9c · 2025-07-30T20:57:44.000Z
diff --git a/pyproject.toml b/pyproject.toml
@@ -101,12 +101,12 @@ jax_core_deps = [
   "protobuf==4.25.5",
 ]
 jax_cpu = [
-  "jax==0.7.0",
+  "jax==0.6.2",
   "algoperf[jax_core_deps]",
 ]
 
 jax_gpu = [
-  "jax[cuda12]==0.7.0",
+  "jax[cuda12]==0.6.2",
   "algoperf[jax_core_deps]",
 ]
 
