fix pyink

[NuojCheng]

Description

Fixes every failure on main that is attributable to recently merged Qwen
explicit-sharding work. Three commits.

1. pyink formatting (Code Quality Check). pre-commit run pyink --all-files
fails on main. Five files landed wrapped at 80 columns instead of the repo's
--pyink-indentation=2 --line-length=122:

• src/maxtext/models/qwen3.py
• src/maxtext/models/qwen3_5.py
• tests/integration/train_tests.py
• tests/unit/pyconfig_test.py
• tests/unit/train_compile_test.py

All five came in with #5119 (f452e71). The Code Quality job runs pyink over all
files rather than only the diff, so this reddens the check on every unrelated PR.
The nightly at 13988fb (run 33932239729) is the first to include #5119 and it
reports exactly these five files. That commit is pyink's output verbatim.

2. test_tpu_qwen3_explicit_sharding_matches_auto tolerance. Fails on
tpu7x-8, passes on v6e-4, bit-identical across nightlies:

Not equal to tolerance rtol=1e-06, atol=0
ACTUAL  (explicit): [189.751099, 189.440186, 189.269714]
DESIRED (auto):     [189.751099, 189.442078, 189.269928]
Max relative difference: 9.99e-06

Only the dense qwen3 subtest is affected -- the one that sets
ici_tensor_parallelism=-1, which resolves to the device count. With
base_num_query_heads=8 that is 4-way TP on v6e (2 heads per shard) and 8-way on
tpu7x (1 head per shard). Step 1 matches bit-for-bit and the drift appears only
from step 2, so the backward reductions reassociate under the pinned layout while
the forward pass is unchanged. Loosened to rtol=1e-4, matching what the sibling
hybrid test already allows for the same effect.

3. test_tpu_qwen3_hybrid_explicit_sharding_matches_auto divisibility. Hard
failure on tpu7x-8, new with #5119:

ValueError: Sharding spec ('tensor',) implies that array axis 2 is partitioned
8 times, but does not evenly divide the dimension size 4.
Got shape: (96, 256, 4, 384)

Same root cause, different symptom: ici_tensor_parallelism=-1 asks the head axis
to shard 8 ways while gdn_num_key_heads=4. It passes on v6e-4 only because 4
divides 4. Pinned the degree at 4 with the leftover devices going to data
parallelism, which both shard modes see alike.

Tests

pre-commit run pyink --all-files and mdformat --all-files pass on this branch;
pyink fails on main. codespell, pylint, yamllint, actionlint and the
decoupled-requirements hook were already passing and still pass.

The qwen3-next fix was reproduced and verified locally by AOT-compiling that exact
config: with ici_tensor_parallelism=-1 it raises the identical ValueError at
v5p-16 (8 devices), and with the pinned degree it compiles clean at both
v5p-8 and v5p-16. The tpu7x CI job on this PR is the end-to-end check for
both test changes.

Not addressed here (not caused by this work)

For the record, the other failures on main at 13988fb were checked and are
unrelated:

• tokamax_test.py::test_smoke_train_tokamax_v1_fp8_full_ep1 -- NotImplementedError: subchannel_iters != 1 not supported yet in tgmm, a tokamax limitation.
• train_tests.py::test_moe_nanoo_fp8[_sparse_matmul] -- GCS 412 PreconditionFailed writing to runner-maxtext-logs, infrastructure.
• sft_multimodal_gemma3_demo.ipynb / native_lora_demo.ipynb -- NOT_FOUND: "google/gemma-3-4b-it" tokenizer, failing nightly for weeks.
• checkpoint_resharding_test.py::test_checkpoint_resharding -- tensorstore zarr3 chunk_shape mismatch. Not root-caused here, but no PR in this range touches checkpointing, orbax or tensorstore, and this shard has a history of orbax-flavoured breakage (e.g. 'Checkpointer' object has no attribute 'restore' in run 33776554492, well before this work).

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. (AOT compile verification described above; tpu7x CI on this PR covers the rest.)
• 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. (Not applicable.)

[gemini-code-assist]

Code Review

This pull request consists entirely of code formatting adjustments across several files, including model definitions (qwen3.py, qwen3_5.py) and test suites (train_tests.py, pyconfig_test.py, train_compile_test.py). Multi-line statements, function calls, and assertions have been consolidated into single lines to improve readability and consistency. No functional changes were introduced, and there are no review comments to address.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!