upgrade_jax

Author: priyakasimbeg

setup.cfg

@@ -121,17 +121,17 @@ jax_core_deps =
 
 # JAX CPU
 jax_cpu =
-  jax==0.4.35
-  jaxlib==0.4.35
+  jax==0.4.38
+  jaxlib==0.4.38
   %(jax_core_deps)s
 
 # JAX GPU
 # Note this installs both jax and jaxlib.
 jax_gpu =
-  jax==0.4.35
-  jaxlib==0.4.35
-  jax-cuda12-plugin[with_cuda]==0.4.35
-  jax-cuda12-pjrt==0.4.35
+  jax==0.4.38
+  jaxlib==0.4.38
+  jax-cuda12-plugin[with_cuda]==0.4.38
+  jax-cuda12-pjrt==0.4.38
   %(jax_core_deps)s
 
 # PyTorch CPU