Fix explicit-mesh loss sharding assert on the NNX path#4208

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

fix/deepseek-batchsplit-context-reshardAI-Hypercomputer/maxtext:fix/deepseek-batchsplit-context-reshardCopy head branch name to clipboard