[NNX] NNX migration prep (4/N): sharding tools, Linen↔NNX checkpoint utilities, and post-training fixes#3652

Draft

ecnal-cienet wants to merge 2 commits intomainfrom

feat/nnx-post-train-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