Add input sharding unit test#3454

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

CIeNET-International:charlesli/input_shardings_unittestCIeNET-International/maxtext:charlesli/input_shardings_unittestCopy head branch name to clipboard