[NNX] NNX migration prep (10/N): vocab tiling custom_vjp with output-head carve-out#3849

Draft

ecnal-cienet wants to merge 10 commits into

mainfrom

feat/nnx-vocab-tiling-custom-vjp

Commits on May 8, 2026

NNX: add TrainState, model creation utilities, and training loop support
ecnal-cienet
committed
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
NNX: AQT in MaxEngine + serve-mode reload + gpt3 prefill fix
ecnal-cienet
committed
NNX: vocab tiling custom_vjp with output-head carve-out
ecnal-cienet
committed