[NNX] NNX migration prep (4.6/N): Linen<->NNX checkpoint comparator

[ecnal-cienet]

NNX Migration Route Map

1. ✅ Add NNX scaffolding: pure_nnx flag, init_state_fn, TrainStateNNX, NNX utils. Linen workflow unchanged. (PR NNX migration prep (1/N): pure_nnx flag and init_state_fn scaffolding #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 NNX migration prep (2/N): NNX utils and sharding utilities #3470)
3. ✅ NNX fully supported end-to-end: TrainStateNNX, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR NNX migration prep (3/N): TrainState, model creation, and end-to-end training loop #3500)
4. ✅ Sharding diagnostics on NNX, plus post-training bugfixes that surfaced once the NNX path got exercised end-to-end. (PR [NNX] NNX migration prep (4/N): sharding tools and post-training fixes #3652)
   4.5. ✅ Linen↔NNX checkpoint converter. (PR [NNX] NNX migration prep (4.5/N): Linen<->NNX checkpoint converter #3843)
   4.6. 🔄 [This PR] Linen↔NNX checkpoint comparator (stacked on PR4.5; the two are file-disjoint).
5. ❌ NNX correctness fixes, feature enablements, and vocab tiling on NNX.
6. ❌ NNX-native DPO.
7. ❌ NNX-native MaxEngine inference.
8. ❌ NNX-native LoRA + GRPO.
9. ❌ NNX-aware QK-Clip + remaining checkpoint utilities.
   9.5. ❌ NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix.
10. ❌ Vocab tiling custom_vjp for NNX.
11. ❌ Set NNX defaults to True; regenerate sharding goldens; flip back integration-test pure_nnx=False annotations.
12. ❌ Delete Linen-specific code paths and NNX compatibility flags.

Description

Companion to PR4.5 (the converter). This PR adds compare_linen_nnx_checkpoint.py — a structural and numerical comparison utility used during the NNX migration to validate parity between Linen and NNX checkpoints. Split out of the original PR4.5 bundle on 2026-05-07 so each PR stays narrowly reviewable. PR4.5 and PR4.6 are file-disjoint, and the comparator does not import the converter — PR4.6 only stacks on PR4.5 so reviewers see the delta cleanly.

This is a pure addition — no existing files are modified, no production-code paths reference the utility, and no Linen or NNX runtime behavior changes. PR5+ do not depend on this branch.

Diff: +1110 / −0 across 2 new files.

What it does

src/maxtext/checkpoint_conversion/compare_linen_nnx_checkpoint.py:

• Three comparison modes — Linen vs NNX (cross-format), Linen vs Linen, NNX vs NNX.
• Format auto-detection — picks Linen vs NNX per checkpoint and applies matching normalization:
  • Strip {value: ...} wrappers on NNX
  • Collapse double params nesting on Linen
  • Filter NNX-only RNG entries
• Cross-format-only transforms — layer-axis transposition between Linen per-layer arrays and NNX stacked tensors only runs for Linen↔NNX comparisons.
• What it reports — tree-structure mismatches, shape mismatches, and (with --compare_values) numerical diffs at configurable --atol / --rtol.

Tests

tests/unit/compare_linen_nnx_checkpoint_test.py — pure-CPU, 60 cases. Covers structural diffing, shape mismatches, numerical comparison at configurable tolerances, and each of the cross- / same-format combinations.

Existing tests untouched.

Stats

• Diff: +1110 / −0 across 2 files (2 new, 0 modified).
• Production-code impact: none. No existing source file imports the utility.
• Linen preservation: trivially preserved — no Linen file is touched.

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 73.73108% with 295 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...ckpoint_conversion/compare_linen_nnx_checkpoint.py	64.61%	99 Missing and 10 partials ⚠️
src/maxtext/trainers/pre_train/train.py	46.51%	71 Missing and 21 partials ⚠️
...xtext/checkpoint_conversion/linen_nnx_converter.py	87.82%	22 Missing and 16 partials ⚠️
src/maxtext/utils/maxtext_utils.py	82.40%	14 Missing and 5 partials ⚠️
src/maxtext/utils/sharding.py	72.22%	6 Missing and 9 partials ⚠️
src/maxtext/trainers/post_train/rl/train_rl.py	62.50%	6 Missing ⚠️
src/maxtext/utils/train_utils.py	79.31%	5 Missing and 1 partial ⚠️
src/maxtext/utils/muon_utils.py	80.95%	4 Missing ⚠️
src/maxtext/utils/model_creation_utils.py	93.87%	2 Missing and 1 partial ⚠️
src/maxtext/common/checkpointing.py	86.66%	1 Missing and 1 partial ⚠️
... and 1 more

📢 Thoughts on this report? Let us know!