[NNX] NNX migration prep (9/N): NNX-aware QK-Clip + checkpoint utilities#3836

Draft

ecnal-cienet wants to merge 7 commits into

mainfrom

feat/nnx-qk-clip-and-checkpoint-utils

Commits on May 12, 2026

NNX migration prep (4/N): sharding tools + post-training bugfixes
ecnal-cienet
committed
NNX migration prep (4.5/N): Linen<->NNX checkpoint converter
ecnal-cienet
committed
NNX: correctness fixes, enable feature paths, and vocab tiling on NNX
ecnal-cienet
committed
NNX: native DPO (TrainStateNNX.reference_model + dpo_loss_fn_nnx)
ecnal-cienet
committed
NNX: native MaxEngine inference (drop route-to-Linen path in maxengine.py)
ecnal-cienet
committed
NNX: native LoRA + GRPO (drop maxengine LoRA carve-out, drop GRPO pure_nnx warning)
ecnal-cienet
committed
NNX: QK-Clip on NNX + NNX-format checkpoint utilities
ecnal-cienet
committed