Truncate out_sharding to match tensor ndim when pspec has extra dimensions. - #4769

Truncate out_sharding to match tensor ndim when pspec has extra dimensions.#4769

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

test_960571888AI-Hypercomputer/maxtext:test_960571888Copy head branch name to clipboard