NNX migration prep (4/N): sharding tools and Linen<->NNX checkpoint utilities

[ecnal-cienet]

NNX Migration Route Map

1. ✅ Add NNX scaffolding: pure_nnx flag, init_state_fn, TrainStateNNX, NNX utils. Linen workflow unchanged. (PR #3427)
2. ✅ NNX sharding utilities: get_abstract_state_nnx, get_named_sharding_nnx, set_named_sharding_nnx, get_partition_spec_nnx, get_mesh_from_config. (PR #3470)
3. ✅ NNX fully supported end-to-end: TrainStateNNX, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR #3500)
4. 🔄 [This PR] NNX sharding diagnostics and bidirectional Linen↔NNX checkpoint conversion utilities.
5. ❌ NNX unit tests and performance verification complete. Set pure_nnx=True as default.
6. ❌ Remove Linen-specific code paths and NNX compatibility flags.

Description

Note: This is the fourth in a series of NNX migration PRs. This PR adds developer tooling to inspect NNX sharding and convert / compare checkpoints across Linen and NNX formats. No training logic is changed.

Sharding diagnostics

• maxtext_utils.py — print_shardings_params now dispatches on pure_nnx: for NNX models it iterates over the flat nnx.State rather than the Linen params tree.
• tests/utils/run_sharding_dump.py — run_single_dump() now propagates --pure_nnx=true to the sharding-dump subprocess when the flag is set, enabling NNX sharding dumps without manual flag threading.

Linen ↔ NNX checkpoint converter

src/maxtext/checkpoint_conversion/linen_nnx_converter.py — a standalone CPU-only script that bidirectionally converts Orbax checkpoints between Linen and NNX formats.

Key transformations handled:

Direction	params tree	opt_state	step	Layer layout
Linen → NNX	params/params/<model> → model/<model> + {value:} wrappers	remove params level from mu/nu	move inside optimizer/	stack layers_N arrays → layers tensor (axis 1)
NNX → Linen	reverse of above	add params level	move to top level	unstack layers tensor → layers_N per-layer arrays

--direction accepts linen_to_nnx, nnx_to_linen, or auto (detects format from checkpoint keys).

Checkpoint comparison utility

src/maxtext/checkpoint_conversion/compare_linen_nnx_checkpoint.py — compares tree structure, shapes, and optionally values between any two Orbax checkpoints (Linen vs NNX, or same-format). Auto-detects format and applies cross-format normalization (layer axis transposition, {value:} unwrapping, RNG filtering) only when needed.

# Structure + shape comparison (Linen vs NNX)
python compare_linen_nnx_checkpoint.py \
  --ckpt_path_1="gs://bucket/linen_checkpoint/0/items" \
  --ckpt_path_2="gs://bucket/nnx_checkpoint/0/items"

# Value comparison
python compare_linen_nnx_checkpoint.py \
  --ckpt_path_1="gs://bucket/ckpt_a/0/items" \
  --ckpt_path_2="gs://bucket/ckpt_b/0/items" \
  --compare_values --atol=1e-5 --rtol=1e-5

Tests

Unit tests:

python3 -m pytest tests/unit/linen_nnx_converter_test.py -v
python3 -m pytest tests/unit/compare_linen_nnx_checkpoint_test.py -v

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.

[codecov]

Codecov Report

❌ Patch coverage is 20.47579% with 936 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xtext/checkpoint_conversion/linen_nnx_converter.py	8.97%	283 Missing and 1 partial ⚠️
...ckpoint_conversion/compare_linen_nnx_checkpoint.py	12.09%	268 Missing and 1 partial ⚠️
src/maxtext/trainers/pre_train/train.py	30.32%	88 Missing and 20 partials ⚠️
src/maxtext/utils/maxtext_utils.py	28.33%	80 Missing and 6 partials ⚠️
src/maxtext/utils/sharding.py	5.55%	50 Missing and 1 partial ⚠️
src/maxtext/utils/train_utils.py	47.91%	21 Missing and 4 partials ⚠️
src/maxtext/utils/maxtext_utils_nnx.py	42.10%	20 Missing and 2 partials ⚠️
src/maxtext/trainers/pre_train/train_compile.py	33.33%	14 Missing and 4 partials ⚠️
src/maxtext/utils/muon_utils.py	14.28%	18 Missing ⚠️
src/maxtext/utils/model_creation_utils.py	52.77%	15 Missing and 2 partials ⚠️
... and 5 more

📢 Thoughts on this report? Let us know!