[NNX] Delete Linen 4/5: remove the Linen decoder/attention layers and *_as_linen model wrappers

[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)
5. ✅ NNX correctness fixes, feature enablements, and vocab tiling on NNX.
6. ✅ NNX-native DPO.
7. ✅ NNX-native MaxEngine inference. (PR [NNX] NNX migration prep (7/N): NNX-native MaxEngine inference #3821)
8. ✅ NNX-native LoRA + GRPO. (PR [NNX] NNX migration prep (8/N): NNX native lora grpo #3824)
9. ✅ NNX-aware QK-Clip + remaining checkpoint utilities. (PR [NNX] NNX migration prep (9/N): NNX-aware QK-Clip + checkpoint utilities #3836)
   9.5. ✅ NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix. (PR [NNX] NNX migration prep (9.5/N): NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix #3844)
10. ✅ Vocab tiling custom_vjp for NNX.
11. ✅ Set NNX defaults to True; regenerate sharding goldens; flip back integration-test pure_nnx=False annotations. (PR [NNX] NNX migration (11/N): set pure_nnx / enable_nnx / pure_nnx_decoder defaults to True #3526)
12. Delete Linen-specific code paths and NNX compatibility flags — delivered as 5 stacked PRs (12.1–12.5):
    12.1. ✅ Collapse Linen dispatch in utils / quantization / checkpointing / sharding. (PR [NNX] Delete Linen 1/5: collapse dispatch in utils, quantization, checkpointing, sharding #4357)
    12.2. ✅ Collapse Linen dispatch in trainers (pre-train, DiLoCo, GRPO); regenerate the NNX DPO golden. (PR [NNX] Delete Linen 2/5: collapse dispatch in trainers (pre-train, DiLoCo, GRPO) #4358)
    12.3. ✅ Collapse Linen dispatch in inference (maxengine, kvcache, vLLM, LoRA). (PR [NNX] Delete Linen 3/5: collapse dispatch in inference (maxengine, kvcache, vLLM, LoRA) #4359)
    12.4. 🔄 [This PR] Delete the Linen model code now that 12.1–12.3 removed every reference: the decoders.py stack (Decoder/DecoderLayer/SequentialBlockDecoderLayers, only deepstack_process kept), the Linen attention/embeddings/encoders/MTP paths, and the *_as_linen/transformer_as_linen wrappers in the model files. The tests importing those classes are removed in this same PR so the branch stays import-clean. The NNX↔Linen bridge + Linen GRPO reference are kept for the checkpoint converters / torch-gated tests (follow-up). Stacks on 12.3.
    12.5. ❌ Remove the pure_nnx / enable_nnx / pure_nnx_decoder config flags.

Description

Fourth of the stacked "Delete Linen" series (see 1/5). PRs 1–3 removed every source reference to the Linen model code, so this PR deletes it. That means the Linen decoder stack in decoders.py (Decoder, DecoderLayer, SequentialBlockDecoderLayers — only deepstack_process stays), the Linen attention implementations, the Linen embeddings / encoders / multi_token_prediction paths, and the *_as_linen / transformer_as_linen ToLinen wrappers in the model files (the wrapped NNX classes stay). The tests that imported those classes are removed in this same PR so the branch stays import-clean — a repo-wide grep for the deleted symbols now returns only the deletions themselves.

Note that the NNX↔Linen bridge (transformer_as_linen / TransformerLinen / init_initial_state) and the Linen GRPO reference are deliberately kept here, since the checkpoint-conversion tools and the torch-gated correctness tests still use them; migrating those to pure NNX is a follow-up. Base branch is feat/nnx-del-linen-3-inference. ~16 files, +56 / −3498.

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

Files with missing lines	Patch %	Lines
src/maxtext/trainers/pre_train/train.py	69.23%	31 Missing and 13 partials ⚠️
src/maxtext/utils/lora_utils.py	0.00%	12 Missing ⚠️
src/maxtext/utils/train_utils.py	83.33%	2 Missing and 1 partial ⚠️
src/maxtext/trainers/pre_train/train_compile.py	87.50%	1 Missing and 1 partial ⚠️
src/maxtext/utils/muon_utils.py	90.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!