new jaxlib dev version

Author: lkolluru05

Dockerfile

@@ -46,12 +46,8 @@ RUN pip install -qq --upgrade pip && \
 FROM base AS ci
 
 # TODO(markblee): Remove gcp,vertexai_tensorboard from CI.
-<<<<<<< HEAD
 RUN uv pip install -qq .[core,audio,orbax,dev,gcp,vertexai_tensorboard] && \
     uv cache clean
-=======
-RUN uv pip install .[core,audio,orbax,dev,gcp,vertexai_tensorboard,open_api] && uv cache clean
->>>>>>> 4109ab0 (Move UV_FIND_LINKS to pyproject.toml)
 COPY . .
 
 # Defaults to an empty string, i.e. run pytest against all files.
@@ -70,11 +66,7 @@ FROM base AS bastion
 # TODO(markblee): Consider copying large directories separately, to cache more aggressively.
 # TODO(markblee): Is there a way to skip the "production" deps?
 COPY . /root/
-<<<<<<< HEAD
 RUN uv pip install -qq .[core,gcp,vertexai_tensorboard] && uv cache clean
-=======
-RUN uv pip install .[core,gcp,vertexai_tensorboard] && uv cache clean
->>>>>>> 4109ab0 (Move UV_FIND_LINKS to pyproject.toml)
 
 ################################################################################
 # Dataflow container spec.                                                     #
@@ -85,11 +77,7 @@ FROM base AS dataflow
 # Beam workers default to creating a new virtual environment on startup. Instead, we want them to
 # pickup the venv setup above. An alternative is to install into the global environment.
 ENV RUN_PYTHON_SDK_IN_DEFAULT_ENVIRONMENT=1
-<<<<<<< HEAD
 RUN uv pip install -qq .[core,gcp,dataflow] && uv cache clean
-=======
-RUN uv pip install .[core,gcp,dataflow] && uv cache clean
->>>>>>> 4109ab0 (Move UV_FIND_LINKS to pyproject.toml)
 COPY . .
 
 # Dataflow workers can't start properly if the entrypoint is not set

axlearn/common/array_serialization.py

@@ -307,6 +307,7 @@ async def _async_serialize(
     )
     # pylint: disable=protected-access
     spec_has_metadata = {
+        "0.6.2.dev0+selfbuilt": lambda: serialization.ts_impl._spec_has_metadata,
         "0.6.2": lambda: serialization.ts_impl._spec_has_metadata,
         "0.5.3": lambda: serialization._spec_has_metadata,
     }[jax.__version__]()
@@ -487,6 +488,7 @@ async def cb(index: array.Index, device: jax.Device):
         requested_domain = ts.IndexTransform(input_shape=shape)[index].domain
         restricted_domain = t.domain.intersect(requested_domain)
         estimate_read_memory_footprint = {
+            "0.6.2.dev0+selfbuilt": lambda: serialization.ts_impl.estimate_read_memory_footprint,
             "0.6.2": lambda: serialization.ts_impl.estimate_read_memory_footprint,
             "0.5.3": lambda: serialization.estimate_read_memory_footprint,
         }[jax.__version__]()
@@ -568,6 +570,7 @@ async def cb(index: array.Index, device: jax.Device):
 
     # pylint: disable=protected-access
     create_async_array_from_callback = {
+        "0.6.2.dev0+selfbuilt": lambda: serialization.ts_impl._create_async_array_from_callback,
         "0.6.2": lambda: serialization.ts_impl._create_async_array_from_callback,
         "0.5.3": lambda: serialization.create_async_array_from_callback,
     }[jax.__version__]()
@@ -653,6 +656,7 @@ def serialize(
         commit_futures = [[] for _ in range(len(tensorstore_specs))]
 
         async_serialize = {
+            "0.6.2.dev0+selfbuilt": lambda: serialization.ts_impl.async_serialize,
             "0.6.2": lambda: serialization.ts_impl.async_serialize,
             "0.5.3": lambda: serialization.async_serialize,
         }[jax.__version__]()

pyproject.toml

@@ -144,8 +144,7 @@ gpu = [
     "nvidia-ml-py>=12.560.30",
     # pin nccl version, otherwise jax[cuda12] will pull latest version
     "nvidia-nccl-cu12==2.27.5",
-    "nvidia-cudnn-cu12>=9.8.0.87" # Pin CuDNN to at least 9.8 for Jax >= 0.6.2
-    "nvidia-cudnn-cu12>=9.8.0.87" # Pin CuDNN to at least 9.8 for Jax >= 0.6.2
+    "nvidia-cudnn-cu12>=9.8.0.87", # Pin CuDNN to at least 9.8 for Jax >= 0.6.2
 ]
 # Open API inference.
 open_api = [