Fix logical to physical axis mapping for attention heads.#4607

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

test_953561420AI-Hypercomputer/maxtext:test_953561420Copy head branch name to clipboard