Upgrade drjax to 0.2.1

Bump the drjax floor to the latest PyPI release in the base requirements
and regenerate the generated_requirements/ locks with the seed-env flow
from docs/development/update_dependencies.md. Regenerating re-resolves
each lock in full, so the locks also pick up unrelated version bumps.

Also pins llguidance<1.8.0 in the post-training base requirements:
seed-env cannot resolve that lock otherwise, since vllm's tpu.txt
constrains llguidance to <1.8.0 while the previously unpinned entry
resolved to >=1.8.0. That break is independent of drjax and reproduces
on unmodified main.

Author: khatwanimohit

src/dependencies/requirements/base_requirements/requirements.txt

@@ -4,7 +4,7 @@ array-record
 chex
 cloud-accelerator-diagnostics
 datasets
-drjax
+drjax>=0.2.1
 flax
 gcsfs
 google-api-python-client

src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt

@@ -12,7 +12,7 @@ gspread
 ipykernel
 ipywidgets
 ijson
-llguidance
+llguidance<1.8.0
 loguru
 lxml
 mistral_common

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

@@ -16,7 +16,7 @@ astroid>=4.0.4
 astunparse>=1.6.3
 attrs>=26.1.0
 auditwheel>=6.8.1
-black>=25.12.0
+black>=26.5.1
 blobfile>=3.3.0
 build>=1.5.0
 certifi>=2026.7.22
@@ -25,14 +25,14 @@ cfgv>=3.5.0
 charset-normalizer>=3.5.1
 cheroot>=11.1.2
 chex>=0.1.92
-click>=8.4.2
+click>=8.5.0
 cloud-accelerator-diagnostics>=0.1.1
 cloudpickle>=3.1.2
 clu>=0.0.12
 colorama>=0.4.6
 contourpy>=1.3.3
-coverage>=7.15.4
-cryptography>=50.0.0
+coverage>=7.16.0
+cryptography>=50.0.1
 cycler>=0.12.1
 datasets>=4.8.5
 decorator>=5.3.1
@@ -42,7 +42,7 @@ distlib>=0.4.3
 distro>=1.9.0
 dm-tree>=0.1.10
 docstring-parser>=0.18.0
-drjax>=0.2.0
+drjax>=0.2.1
 editdistance>=0.8.1
 einops>=0.8.2
 einshape>=1.0
@@ -61,25 +61,25 @@ gast>=0.7.0
 gcsfs>=2026.2.0
 google-api-core>=2.34.0
 google-api-python-client>=2.199.0
-google-auth>=2.56.3
-google-auth-httplib2>=0.4.1
-google-auth-oauthlib>=1.4.0
-google-cloud-aiplatform>=1.165.1
+google-auth>=2.57.0
+google-auth-httplib2>=0.4.2
+google-auth-oauthlib>=1.4.1
+google-cloud-aiplatform>=2.0.1
 google-cloud-appengine-logging>=1.10.0
-google-cloud-audit-log>=0.6.1
-google-cloud-bigquery>=3.43.0
-google-cloud-core>=2.6.1
-google-cloud-logging>=3.16.2
+google-cloud-audit-log>=0.6.2
+google-cloud-bigquery>=3.44.0
+google-cloud-core>=2.7.0
+google-cloud-logging>=3.16.3
 google-cloud-mldiagnostics>=1.0.7
 google-cloud-monitoring>=2.31.0
 google-cloud-resource-manager>=1.18.0
 google-cloud-storage>=3.13.1
-google-cloud-storage-control>=1.13.0
+google-cloud-storage-control>=1.14.0
 google-crc32c>=1.8.0
-google-genai>=2.19.0
+google-genai>=2.20.0
 google-pasta>=0.2.0
-google-resumable-media>=2.10.1
-googleapis-common-protos>=1.75.1
+google-resumable-media>=2.10.2
+googleapis-common-protos>=1.75.2
 grain>=0.2.18
 grpc-google-iam-v1>=0.14.5
 grpcio>=1.83.0
@@ -91,24 +91,24 @@ hf-xet>=1.6.0 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or
 httpcore>=1.0.9
 httplib2>=0.32.0
 httpx>=0.28.1
-huggingface-hub>=1.28.0
+huggingface-hub>=1.29.0
 humanize>=4.16.0
 hypothesis>=6.165.10
 identify>=2.6.19
 idna>=3.18
 immutabledict>=4.3.1
 importlab>=0.8.1
-importlib-metadata>=9.0.0
+importlib-metadata>=9.0.1
 iniconfig>=2.3.0
-isort>=8.0.1
+isort>=9.0.1
 jaraco-functools>=4.6.0
 jax>=0.11.1
 jax-cuda12-pjrt>=0.11.1 ; sys_platform == 'linux'
 jax-cuda12-plugin>=0.11.1 ; sys_platform == 'linux'
 jaxlib>=0.11.1
 jaxtyping>=0.3.11
 jinja2>=3.1.6
-joblib>=1.5.3
+joblib>=1.6.0
 jsonlines>=4.0.0
 keras>=3.15.1
 kiwisolver>=1.5.0
@@ -128,15 +128,15 @@ ml-dtypes>=0.6.0
 ml-goodput-measurement>=0.2.3
 more-itertools>=11.1.0
 mpmath>=1.4.1
-msgpack>=1.2.1
+msgpack>=1.2.2
 msgspec>=0.21.1
 multidict>=6.7.1
 multiprocess>=0.70.19
 mypy-extensions>=1.1.0
 namex>=0.1.0
 nest-asyncio>=1.6.0 ; sys_platform == 'win32'
 networkx>=3.6.1
-ninja>=1.13.0
+ninja>=1.13.2
 nltk>=3.10.3
 nodeenv>=1.10.0
 numpy>=2.1.3
@@ -171,14 +171,14 @@ parameterized>=0.9.0
 pathspec>=1.1.1
 pathwaysutils>=0.1.11
 pillow>=12.3.0
-platformdirs>=4.11.4
+platformdirs>=4.11.5
 pluggy>=1.6.0
 portpicker>=1.6.0
 pre-commit>=4.6.2
 prometheus-client>=0.26.0
 promise>=2.3
 propcache>=0.5.2
-proto-plus>=1.28.3
+proto-plus>=1.28.4
 protobuf>=6.33.6
 psutil>=7.2.2
 pyarrow>=25.0.1
@@ -187,26 +187,26 @@ pyasn1-modules>=0.4.2
 pycnite>=2024.7.31
 pycparser>=3.0 ; implementation_name != 'PyPy' and platform_python_implementation != 'PyPy'
 pycryptodomex>=3.23.0
-pydantic>=2.13.4
-pydantic-core>=2.46.4
+pydantic>=2.13.5
+pydantic-core>=2.46.5
 pydot>=4.0.1
 pyelftools>=0.33
 pyglove>=0.4.5
 pygments>=2.21.0
-pyink>=25.12.0
-pylint>=4.0.7
+pyink>=26.5.1
+pylint>=4.0.8
 pynvml>=13.0.1
 pyparsing>=3.3.2
 pyproject-hooks>=1.2.0
 pytest>=9.1.1
 pytest-xdist>=3.8.0
 python-dateutil>=2.9.0.post0
-python-discovery>=1.5.3
+python-discovery>=1.6.0
 pytokens>=0.4.1
 pytype>=2024.10.11
 pyyaml>=6.0.3
 qwix>=0.1.8
-regex>=2026.7.19
+regex>=2026.8.31
 requests>=2.34.2
 requests-oauthlib>=2.0.0
 responses>=0.18.0
@@ -220,7 +220,7 @@ setuptools>=84.0.0
 shellingham>=1.5.4
 shortuuid>=1.0.13
 simple-parsing>=0.1.9
-simplejson>=4.1.1
+simplejson>=4.1.2
 six>=1.17.0
 sniffio>=1.3.1
 sortedcontainers>=2.4.0
@@ -239,32 +239,33 @@ tensorflow-text>=2.20.1
 tensorstore>=0.1.85
 termcolor>=3.3.0
 tiktoken>=0.14.0
-tokamax>=0.0.12
-tokenizers>=0.22.2
+tokamax>=0.0.13
+tokenizers>=0.23.1
 toml>=0.10.2
 tomlkit>=0.15.1
 toolz>=1.1.0
 tqdm>=4.70.0
 transformer-engine>=2.18.0
 transformer-engine-cu12>=2.18.0
 transformer-engine-jax>=2.18.0
-transformers>=5.15.1
+transformers>=5.16.1
 treescope>=0.1.10
 typeguard>=2.13.3
-typer>=0.27.1
+typer>=0.27.2
 typing-extensions>=4.16.0
 typing-inspection>=0.4.4
 tzdata>=2026.3 ; sys_platform == 'emscripten' or sys_platform == 'win32'
 uritemplate>=4.2.0
 urllib3>=2.7.0
 uvicorn>=0.52.4
 uvloop>=0.22.1
-virtualenv>=21.7.4
+virtualenv>=21.7.7
 wadler-lindig>=0.1.7
 websockets>=16.1.1
 werkzeug>=3.1.8
 wheel>=0.48.0
-wrapt>=2.3.0
+wrapt>=2.4.0
+xprof>=2.20.0
 xxhash>=4.0.1
 yarl>=1.24.5
 zipp>=4.1.0