Update JAX to 0.10.0 for pre-training

Author: SurbhiJainUSC

docs/development/update_dependencies.md

@@ -78,10 +78,10 @@ Next, run `generate_requirements.sh` to generate the new requirements files. Thi
 script wraps the `seed-env` CLI and handles exporting the lock, and applying any
 overrides. You will need to do this separately for the TPU and GPU environments.
 
-> **Note:** The current `src/dependencies/requirements/generated_requirements/` in the repository were generated using JAX build commit hash: [e0d2967b50abbefd651d563dbcd7afbcb963d08c](https://github.com/jax-ml/jax/commit/e0d2967b50abbefd651d563dbcd7afbcb963d08c).
-
 ### TPU Pre-Training
 
+> **Note:** The current `src/dependencies/requirements/generated_requirements/tpu-requirements.txt` in the repository was generated using JAX build commit hash: efd6cf797ee9c4f29c6c6d5e91ae4432209063be. When regenerating the requirements, either use the same commit hash or update this hash if you use a different one.
+
 If you have made changes to TPU pre-training dependencies in `src/dependencies/requirements/base_requirements/tpu-requirements.txt`, you need to regenerate the pinned pre-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
 
 ```bash
@@ -98,6 +98,8 @@ mv generated_artifacts/python3_12/tpu-requirements.txt \
 
 ### TPU Post-Training
 
+> **Note:** The current `src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt` in the repository was generated using JAX build commit hash: e0d2967b50abbefd651d563dbcd7afbcb963d08c. When regenerating the requirements, either use the same commit hash or update this hash if you use a different one.
+
 If you have made changes to the post-training dependencies in `src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt`, you need to regenerate the pinned post-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
 
 ```bash
@@ -115,6 +117,8 @@ mv generated_artifacts/python3_12/tpu-post-train-requirements.txt \
 
 ### GPU Pre-Training
 
+> **Note:** The current `src/dependencies/requirements/generated_requirements/cuda12-requirements.txt` in the repository was generated using JAX build commit hash: efd6cf797ee9c4f29c6c6d5e91ae4432209063be. When regenerating the requirements, either use the same commit hash or update this hash if you use a different one.
+
 If you have made changes to the GPU pre-training dependencies in `src/dependencies/requirements/base_requirements/cuda12-requirements.txt`, you need to regenerate the pinned pre-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
 
 ```bash

src/dependencies/extra_deps/tpu_overrides.txt

@@ -1 +1,2 @@
-libtpu>=0.0.38
+datasets>=4.8.5
+fsspec>=2023.1.0,<=2026.2.0

src/dependencies/requirements/generated_requirements/cuda12-requirements.txt

@@ -14,14 +14,14 @@ aqtp>=0.9.0
 array-record>=0.8.3
 astroid>=4.0.4
 astunparse>=1.6.3
-attrs>=25.4.0
+attrs>=26.1.0
 auditwheel>=6.6.0
 black>=25.12.0
-build>=1.4.0
+build>=1.4.3
 certifi>=2026.2.25
 cffi>=2.0.0 ; platform_python_implementation != 'PyPy'
 cfgv>=3.5.0
-charset-normalizer>=3.4.6
+charset-normalizer>=3.4.7
 chex>=0.1.91
 click>=8.3.3
 cloud-accelerator-diagnostics>=0.1.1
@@ -30,13 +30,11 @@ cloudpickle>=3.1.2
 clu>=0.0.12
 colorama>=0.4.6
 contourpy>=1.3.3
-cryptography>=47.0.0
+cryptography>=48.0.0
 cycler>=0.12.1
-dataclasses-json>=0.6.7
-datasets>=4.8.5
+datasets>=2.14.4
 decorator>=5.2.1
-deprecated>=1.3.1
-dill>=0.4.1
+dill>=0.3.7
 distlib>=0.4.0
 distro>=1.9.0
 dm-tree>=0.1.10
@@ -48,72 +46,71 @@ einshape>=1.0
 etils>=1.14.0
 execnet>=2.1.2
 fastapi>=0.136.1
-filelock>=3.20.3
+filelock>=3.28.0
 flatbuffers>=25.12.19
-flax>=0.12.6
+flax>=0.12.7
 fonttools>=4.62.1
 frozenlist>=1.8.0
-fsspec>=2026.2.0
+fsspec>=2026.3.0
 gast>=0.7.0
-gcsfs>=2026.2.0
+gcsfs>=2026.5.0
 google-api-core>=2.30.3
-google-api-python-client>=2.194.0
-google-auth>=2.49.2
-google-auth-httplib2>=0.3.1
-google-auth-oauthlib>=1.3.1
-google-cloud-aiplatform>=1.148.1
+google-api-python-client>=2.196.0
+google-auth>=2.52.0
+google-auth-httplib2>=0.4.0
+google-auth-oauthlib>=1.4.0
+google-cloud-aiplatform>=1.151.0
 google-cloud-appengine-logging>=1.9.0
 google-cloud-audit-log>=0.5.0
 google-cloud-bigquery>=3.41.0
-google-cloud-core>=2.5.1
+google-cloud-core>=2.6.0
 google-cloud-logging>=3.15.0
 google-cloud-mldiagnostics>=1.0.2
 google-cloud-monitoring>=2.30.0
 google-cloud-resource-manager>=1.17.0
 google-cloud-storage>=3.10.1
 google-cloud-storage-control>=1.11.0
 google-crc32c>=1.8.0
-google-genai>=1.73.1
+google-genai>=1.75.0
 google-pasta>=0.2.0
-google-resumable-media>=2.8.2
-googleapis-common-protos>=1.74.0
+google-resumable-media>=2.9.0
+googleapis-common-protos>=1.75.0
 grain>=0.2.16
 grpc-google-iam-v1>=0.14.4
-grpcio>=1.78.0
-grpcio-status>=1.78.0
+grpcio>=1.80.0
+grpcio-status>=1.80.0
 gviz-api>=1.10.0
 h11>=0.16.0
 h5py>=3.14.0
-hf-xet>=1.4.3 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
+hf-xet>=1.5.0 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
 httpcore>=1.0.9
 httplib2>=0.31.2
 httpx>=0.28.1
-huggingface-hub>=1.12.0
+huggingface-hub>=0.36.2
 humanize>=4.15.0
 hypothesis>=6.142.1
 identify>=2.6.19
 idna>=3.11
 immutabledict>=4.3.1
 importlab>=0.8.1
-importlib-metadata>=9.0.0
+importlib-metadata>=8.7.1
 iniconfig>=2.3.0
 isort>=8.0.1
-jax>=0.9.2
-jax-cuda12-pjrt>=0.9.2 ; sys_platform == 'linux'
-jax-cuda12-plugin>=0.9.2 ; sys_platform == 'linux'
-jaxlib>=0.9.2
+jax>=0.10.0
+jax-cuda12-pjrt>=0.10.0 ; sys_platform == 'linux'
+jax-cuda12-plugin>=0.10.0 ; sys_platform == 'linux'
+jaxlib>=0.10.0
 jaxtyping>=0.3.9
 jinja2>=3.1.6
 jsonlines>=4.0.0
-keras>=3.13.2
+keras>=3.14.0
 kiwisolver>=1.5.0
 latex2sympy2-extended>=1.11.0
 libclang>=18.1.1
 libcst>=1.8.6
 markdown>=3.10.2
 markdown-it-py>=4.0.0
 markupsafe>=3.0.3
-marshmallow>=3.26.2
 math-verify>=0.9.0
 matplotlib>=3.10.8
 mccabe>=0.7.0
@@ -125,7 +122,7 @@ mpmath>=1.3.0
 msgpack>=1.1.2
 msgspec>=0.21.1
 multidict>=6.7.1
-multiprocess>=0.70.19
+multiprocess>=0.70.15
 mypy-extensions>=1.1.0
 namex>=0.1.0
 nest-asyncio>=1.6.0 ; sys_platform == 'win32'
@@ -134,77 +131,76 @@ ninja>=1.13.0
 nodeenv>=1.10.0
 numpy>=2.0.2
 numpy-typing-compat>=20251206.2.0
-nvidia-cublas-cu12>=12.9.1.4 ; sys_platform == 'linux'
-nvidia-cuda-cccl>=13.2.27
+nvidia-cublas-cu12>=12.9.2.10 ; sys_platform == 'linux'
+nvidia-cuda-cccl>=13.2.75
 nvidia-cuda-cccl-cu12>=12.9.27
 nvidia-cuda-cupti-cu12>=12.9.79 ; sys_platform == 'linux'
 nvidia-cuda-nvcc-cu12>=12.9.86 ; sys_platform == 'linux'
 nvidia-cuda-nvrtc-cu12>=12.9.86 ; sys_platform == 'linux'
 nvidia-cuda-runtime-cu12>=12.9.79 ; sys_platform == 'linux'
-nvidia-cudnn-cu12>=9.20.0.48 ; sys_platform == 'linux'
+nvidia-cudnn-cu12>=9.21.0.82 ; sys_platform == 'linux'
 nvidia-cufft-cu12>=11.4.1.4 ; sys_platform == 'linux'
 nvidia-cusolver-cu12>=11.7.5.82 ; sys_platform == 'linux'
 nvidia-cusparse-cu12>=12.5.10.65 ; sys_platform == 'linux'
 nvidia-nccl-cu12>=2.29.7 ; sys_platform == 'linux'
 nvidia-nvjitlink-cu12>=12.9.86 ; sys_platform == 'linux'
-nvidia-nvshmem-cu12>=3.5.21 ; sys_platform == 'linux'
+nvidia-nvshmem-cu12>=3.6.5 ; sys_platform == 'linux'
 oauthlib>=3.3.1
 omegaconf>=2.3.0
-opentelemetry-api>=1.16.0
+opentelemetry-api>=1.41.1
 opt-einsum>=3.4.0
 optax>=0.2.8
 optree>=0.19.0
 optype>=0.17.0
-orbax-checkpoint>=0.11.36
-orbax-export>=0.0.8
-packaging>=26.0
+orbax-checkpoint>=0.11.39
+packaging>=26.1
 pandas>=3.0.2
 parameterized>=0.9.0
 pathspec>=1.1.1
 pathwaysutils>=0.1.8
-pillow>=12.1.1
+pillow>=12.2.0
 platformdirs>=4.9.6
 pluggy>=1.6.0
 portpicker>=1.6.0
 pre-commit>=4.6.0
 promise>=2.3
 propcache>=0.4.1
-proto-plus>=1.27.2
+proto-plus>=1.28.0
 protobuf>=6.33.6
 psutil>=7.2.2
 pyarrow>=24.0.0
 pyasn1>=0.6.3
 pyasn1-modules>=0.4.2
 pycnite>=2024.7.31
 pycparser>=3.0 ; implementation_name != 'PyPy' and platform_python_implementation != 'PyPy'
-pydantic>=2.13.3
-pydantic-core>=2.46.3
+pydantic>=2.13.4
+pydantic-core>=2.46.4
 pydot>=4.0.1
 pyelftools>=0.32
 pyglove>=0.4.5
-pygments>=2.19.2
+pygments>=2.20.0
 pyink>=25.12.0
 pylint>=4.0.5
 pyparsing>=3.3.2
 pyproject-hooks>=1.2.0
 pytest>=8.4.2
 pytest-xdist>=3.8.0
 python-dateutil>=2.9.0.post0
+python-discovery>=1.3.0
 pytokens>=0.4.1
 pytype>=2024.10.11
 pyyaml>=6.0.3
 qwix>=0.1.6
 regex>=2026.4.4
-requests>=2.32.5
+requests>=2.33.1
 requests-oauthlib>=2.0.0
-rich>=14.3.3
+rich>=15.0.0
 safetensors>=0.7.0
 scipy>=1.17.1
-scipy-stubs>=1.17.1.2
+scipy-stubs>=1.17.1.4
 sentencepiece>=0.2.1
 seqio>=0.0.20
 setuptools>=82.0.1
-shellingham>=1.5.4
 simple-parsing>=0.1.8
 simplejson>=4.1.1
 six>=1.17.0
@@ -219,7 +215,7 @@ tensorboard-data-server>=0.7.2
 tensorboard-plugin-profile>=2.13.0
 tensorboardx>=2.6.5
 tensorflow>=2.20.0
-tensorflow-datasets>=4.9.9
+tensorflow-datasets>=4.9.10
 tensorflow-metadata>=1.17.3
 tensorflow-text>=2.20.1
 tensorstore>=0.1.82
@@ -231,28 +227,26 @@ toml>=0.10.2
 tomlkit>=0.14.0
 toolz>=1.1.0
 tqdm>=4.67.3
-transformer-engine>=2.14.0
-transformer-engine-cu12>=2.14.0
-transformer-engine-jax>=2.14.0
-transformers>=5.6.2
+transformer-engine>=2.14.1
+transformer-engine-cu12>=2.14.1
+transformer-engine-jax>=2.14.1
+transformers>=4.57.6
 treescope>=0.1.10
 typeguard>=2.13.3
-typer>=0.25.0
 typing-extensions>=4.15.0
-typing-inspect>=0.9.0
 typing-inspection>=0.4.2
 tzdata>=2026.2 ; sys_platform == 'emscripten' or sys_platform == 'win32'
 uritemplate>=4.2.0
 urllib3>=2.6.3
 uvicorn>=0.46.0
 uvloop>=0.22.1
-virtualenv>=20.36.1
+virtualenv>=21.3.1
 wadler-lindig>=0.1.7
 websockets>=16.0
 werkzeug>=3.1.8
 wheel>=0.46.3
 wrapt>=2.1.2
 xxhash>=3.7.0
 yarl>=1.23.0
-zipp>=3.23.0
+zipp>=3.23.1
 zstandard>=0.25.0