Support packing for TPU Ulysses attention

[huytransformer]

Description

Corollary follow up of PR #4687.

This PR enables sequence packing when using TPU Ulysses attention by lifting the config ban and adding tests.

Tests

• tests/unit/attention_test.py -k ulysses on TPU v4: 7 passed
• tests/unit/configs_value_test.py: Passed.
• Correctness: 20 training steps vs all-gather, c4 data, cp=4, packing=true. Max loss difference is 0.001.

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 tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request enables sequence packing support for TPU Ulysses context parallelism. It removes the validation checks that previously restricted packing with Ulysses, updates the documentation accordingly, and adds unit tests to verify Ulysses attention works correctly with packed sequences. There are no review comments to address, so I have no additional feedback to provide.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/train_utils.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @huytransformer, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This pull request introduces support for sequence packing under TPU Ulysses context parallelism in MaxText. The changes are highly focused, robustly implemented, and include both configuration validation updates and comprehensive mathematical equivalence and gradient equivalence tests.

🔍 General Feedback

• High Code Quality: The implementation is exceptionally clean, with precise removal of the previous constraint and clear updates to the relevant config and train helper files.
• Strong Test Coverage: The tests are well-parameterized and verify both forward output correctness and backward gradient correctness under packed TPU Ulysses configurations.
• Precise Documentation: The documentation in docs/guides/optimization/sharding.md was accurately updated to reflect the new capabilities.