Fix double-compilation in train_step by matching input sharding.#4174

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

igorts/b440499679-train_step-hloAI-Hypercomputer/maxtext:igorts/b440499679-train_step-hloCopy head branch name to clipboard