upgrade jax to 0.5.3

priyakasimbeg · priyakasimbeg · commit 75d63157b8a7 · 2025-04-01T22:33:26.000Z
diff --git a/pyproject.toml b/pyproject.toml
@@ -106,11 +106,11 @@ jax_core_deps = [
   "protobuf==4.25.5",
 ]
 jax_cpu = [
-  "jax==0.5.2",
+  "jax==0.5.3",
   "algoperf[jax_core_deps]",
 ]
 jax_gpu = [
-  "jax[cuda12]==0.5.2",
+  "jax[cuda12]==0.5.3",
   "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==0.5.2",`
	`109`	`+ "jax==0.5.3",`
`110`	`110`	`"algoperf[jax_core_deps]",`
`111`	`111`	`]`
`112`	`112`	`jax_gpu = [`
`113`		`- "jax[cuda12]==0.5.2",`
	`113`	`+ "jax[cuda12]==0.5.3",`
`114`	`114`	`"algoperf[jax_core_deps]",`
`115`	`115`	`]`
`116`	`116`	`pytorch_cpu = ["torch==2.5.1", "torchvision==0.20.1"]`