feat(models): add MaxText pipeline parallelism configuration

[sontungkieu]

Summary

• add a validated MaxTextPipelineConfig for hybrid pipeline/tensor layouts
• route MaxText-specific mesh and runtime overrides through AutoModel
• document PP2xTP4 and PP4xTP2 usage
• cover mesh axes, batch divisibility, and invalid layouts with regression tests

Motivation

Tunix MaxText users currently need to hand-assemble the full MaxText mesh and pipeline override set. This adds a reusable typed configuration while leaving ordinary MaxText and multi-slice configuration paths unchanged.

Validation

• 10 passed in tests/models/maxtext_parallelism_test.py on CPU JAX
• pyink/isort checks on the touched MaxText integration files
• runtime evidence from an 8-device TPU v5e benchmark exists downstream; this draft PR does not claim training or decode throughput

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.