NNX migration prep (2/N): NNX utils and sharding utilities#3470

copybara-service[bot] merged 1 commit intomainAI-Hypercomputer/maxtext:mainfrom

feat/migrate-nnx-utilsAI-Hypercomputer/maxtext:feat/migrate-nnx-utilsCopy head branch name to clipboard