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

[copybara-service]

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

Imported from GitHub PR #4360

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.

Copybara import of the project:

--
642e6d0 by Lance Wang lance.wang@cienet.com:

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

Merging this change closes #4360

Reverts f8d1f3b

[codecov]

Codecov Report

❌ Patch coverage is 39.47368% with 506 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/decoders.py	32.97%	371 Missing and 66 partials ⚠️
src/maxtext/models/models.py	63.96%	26 Missing and 14 partials ⚠️
src/maxtext/layers/embeddings.py	60.00%	8 Missing ⚠️
src/maxtext/layers/encoders.py	50.00%	4 Missing ⚠️
src/maxtext/models/gemma3.py	50.00%	4 Missing ⚠️
src/maxtext/models/qwen3.py	60.00%	4 Missing ⚠️
src/maxtext/layers/nnx_decoders.py	50.00%	2 Missing ⚠️
src/maxtext/models/llama4.py	50.00%	2 Missing ⚠️
src/maxtext/layers/attention_mla.py	75.00%	1 Missing ⚠️
src/maxtext/layers/attention_op.py	66.66%	1 Missing ⚠️
... and 3 more

📢 Thoughts on this report? Let us know!