[NNX] NNX migration prep (5/N): correctness fixes and feature enablements#3766

Draft

ecnal-cienet wants to merge 3 commits intomainfrom

feat/nnx-correctness-fixes

Commits on Apr 29, 2026

NNX: add TrainState, model creation utilities, and training loop support
ecnal-cienet
committed
NNX: add sharding tools, Linen<->NNX checkpoint utilities, and post-training fixes
ecnal-cienet
committed
NNX: correctness fixes, enable feature paths, and vocab tiling on NNX
ecnal-cienet
committed