set jax to 0.5.1

priyakasimbeg · priyakasimbeg · commit 52bb87510d7c · 2025-03-20T21:18:42.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"]
diff --git a/submission_runner.py b/submission_runner.py
@@ -64,6 +64,7 @@
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # Disables tensorRT, cuda warnings.
 # disable only for deepspeech if it works fine for other workloads
 os.environ['XLA_FLAGS'] = '--xla_gpu_enable_triton_gemm=false'
+os.environ['JAX_TRACEBACK_FILTERING'] = "off"
 
 # TODO(znado): make a nicer registry of workloads that lookup in.
 BASE_WORKLOADS_DIR = workloads.BASE_WORKLOADS_DIR

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"]`