chore: Upgrade MaxText post-training pinned dependencies to JAX 0.11.1

[darisoy]

Description

Upgrades MaxText TPU post-training pinned dependencies to JAX 0.11.1 ecosystem:

• jax==0.11.1
• jaxlib==0.11.1
• libtpu==0.0.46
• flax==0.12.9
• optax==0.2.8
• torch==2.11.0+cpu and torchvision==0.26.0+cpu
• Added required resolution overrides (mpmath<=1.3.0,>=1.3.0, llguidance<1.8.0,>=1.7.0, boto3>=1.34.0, tensorflow==2.20.0, tensorflow-text==2.20.1, tokamax>=0.0.13, xprof<=2.23.0).
• Updated generate_requirements.sh post-training resolution flow to inject override-dependencies and Linux x86_64 target platform constraints.
• Updated MaxTextCheckpointManager in distillation_utils.py to correctly initialize and delegate to Tunix's CheckpointManager with MaxText GrainCheckpointHandler.

FIXES: b/552616887

cc @bvandermoon @SurbhiJainUSC

Tests

• Unit tests in tests/post_training/unit/distillation_checkpointing_test.py (2 passed)
• Unit tests in tests/post_training/unit/train_distill_test.py (23 passed, 1 skipped)
• Unit tests in tests/post_training/unit/dpo_trainer_correctness_test.py (2 passed)
• Full post-training unit suite tests/post_training/unit/ (56 passed, 271 skipped)

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have added or updated tests accordingly.
• I have updated documentation accordingly.
• I have included unit tests for new features or bug fixes.
• I have added unit test and integration test coverage for all modified code paths.

[gemini-code-assist]

Code Review

This pull request updates the dependency generation scripts to support environment-specific constraints and overrides, updates post-training TPU dependencies, and refactors MaxTextCheckpointManager to use Orbax's CheckpointManager directly for saving and restoring model parameters, optimizer states, and input pipeline iterators. The review feedback suggests using uv run --with toml in generate_requirements.sh to ensure the toml library is available when injecting overrides into pyproject.toml without requiring manual installation.

[gemini-code-assist]

The script uses python3 -c "import toml ..." to inject overrides into pyproject.toml. However, the toml library is not part of the Python standard library and may not be installed in the host or active virtual environment, which would cause the script to fail with a ModuleNotFoundError.

Since uv is already required and used throughout this script, you can use uv run --with toml python3 to execute the Python script. This guarantees that toml is available in a transient environment without requiring manual installation by the developer.

Suggested change

	python3 -c "
	uv run --with toml python3 -c "

[darisoy]

Addressed in commit fdd06ea using uv run --no-project --with toml python3 to execute the Python script in an isolated ephemeral environment without host dependency requirements.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/moe.py	0.00%	6 Missing ⚠️

📢 Thoughts on this report? Let us know!