set jax to 0.5.1

priyakasimbeg · priyakasimbeg · commit 2cfa2a99bd15 · 2025-03-20T21:20:18.000Z
diff --git a/pyproject.toml b/pyproject.toml
@@ -106,11 +106,11 @@ jax_core_deps = [
   "protobuf==4.25.5",
 ]
 jax_cpu = [
-  "jax",
+  "jax==0.5.2",
   "algoperf[jax_core_deps]",
 ]
 jax_gpu = [
-  "jax[cuda12]",
+  "jax[cuda12]==0.5.2",
   "algoperf[jax_core_deps]",
 ]
 pytorch_cpu = ["torch==2.5.1", "torchvision==0.20.1"]

Original file line number	Diff line number	Diff line change
`@@ -106,11 +106,11 @@ jax_core_deps = [`
`106`	`106`	`"protobuf==4.25.5",`
`107`	`107`	`]`
`108`	`108`	`jax_cpu = [`
`109`		`- "jax",`
	`109`	`+ "jax==0.5.2",`
`110`	`110`	`"algoperf[jax_core_deps]",`
`111`	`111`	`]`
`112`	`112`	`jax_gpu = [`
`113`		`- "jax[cuda12]",`
	`113`	`+ "jax[cuda12]==0.5.2",`
`114`	`114`	`"algoperf[jax_core_deps]",`
`115`	`115`	`]`
`116`	`116`	`pytorch_cpu = ["torch==2.5.1", "torchvision==0.20.1"]`