Update Jax to 0.8.3

GitOrigin-RevId: 792d249

Author: andersensam

axlearn/common/utils_spmd.py

@@ -95,28 +95,5 @@ def setup(
                 # local_device_ids arg allows us to maintain expected behavior
                 init_kwargs["local_device_ids"] = list(range(8))
 
-        # pylint: disable-next=import-outside-toplevel
-        from jax._src.lib import _jax
-
-        if jax.__version__ == "0.8.2":
-            # Jax-0.8.2 introduced a bug in  https://github.com/openxla/xla/pull/35298
-            # Client only sends one heartbeat within heartbeat_timeout,
-            # so server might not receive a heartbeat
-            # in time, and would consider the client dead.
-            # The bug is fixed in https://github.com/openxla/xla/pull/35589 but not in jax-0.8.2
-            # We patch the fix here, so client sends heartbeats twice within heartbeat_timeout.
-            # TODO(ethanli): remove this hack.
-            old = _jax.get_distributed_runtime_client
-
-            def f(*args, **kwargs):
-                logging.info(
-                    "Reducing client heartbeat_timeout %s by half for a bug fix in jax-0.8.2",
-                    kwargs["heartbeat_timeout"],
-                )
-                kwargs["heartbeat_timeout"] = kwargs["heartbeat_timeout"] // 2
-                return old(*args, **kwargs)
-
-            _jax.get_distributed_runtime_client = f
-
         jax.distributed.initialize(**init_kwargs)
         _jax_distributed_initialized = True

pyproject.toml

@@ -23,8 +23,8 @@ core = [
     "absl-py==2.3.1",
     "chex==0.1.91",
     "importlab==0.8.1",  # breaks pytype on 0.8
-    "jax==0.8.2",
-    "jaxlib==0.8.2",
+    "jax==0.8.3",
+    "jaxlib==0.8.3",
     "ml-dtypes==0.5.3",
     "msgpack==1.1.0",  # for checkpointing.
     "nltk==3.7",  # for text preprocessing
@@ -105,7 +105,7 @@ gcp = [
 # Note: Specify -f https://storage.googleapis.com/jax-releases/libtpu_releases.html during install.
 tpu = [
     "axlearn[gcp]",
-    "jax[tpu]==0.8.2",  # must be >=0.4.19 for compat with v5p.
+    "jax[tpu]==0.8.3",  # must be >=0.4.19 for compat with v5p.
     "pathwaysutils==0.1.2",  # For JAX+Pathways single-controller accelerator coordinator.
 ]
 # Vertex AI tensorboard. TODO(markblee): Merge with `gcp`.
@@ -131,7 +131,7 @@ dataflow = [
 # GPU custom kernel dependency.
 gpu = [
     "triton>=2.1.0",
-    "jax[cuda12]==0.8.2",
+    "jax[cuda12]==0.8.3",
     "nvidia-ml-py==12.560.30",
     # pin nccl version, otherwise jax[cuda12] will pull latest version
     "nvidia-nccl-cu12==2.27.5",